[BugFix] Fix H100 CuTeDSL full-run edge cases

[JayceSu98]

Context

This PR collects the TileKernels-side fixes needed after running the full H100 CuTeDSL correctness and benchmark coverage with a local TileLang CuTeDSL backend.

The failures were not tied to a single benchmark configuration. They came from broader edge cases that full coverage exposes more reliably:

• SM90 persistent-block counts that are valid for the supported hardware but are not divisible by four.
• MoE fused layout paths where kernel loops are rounded up to warp/alignment extents while the logical tensor shape is smaller.
• Quantization edge tiles where scaling-factor grids or packed E5M6 output rows/columns are smaller than the scheduled tile.
• BF16/narrow-input per-token quantization paths where the absmax accumulator should be FP32 rather than the input storage dtype.

Changes

Engram

• Replace the fixed num_batches = 4 assumption in engram_grad_w_reduce_kernel.py with the largest batch count up to four that divides the persistent block count.
• This keeps the original layout when possible, but allows SM90 GPUs such as H100 PCIe, which exposes 114 SMs, to run the kernel instead of failing an assertion before code generation.

MoE

• Guard expand_to_fused writes with the real hidden and scaling-factor extents when iterating over aligned loop sizes.
• Initialize the shared cumsum buffer for all threads in fused mapping.
• Avoid speculative out-of-range topk_idx loads in the aligned tail of get_fused_mapping.
• Reserve a rounded-up upper bound for independently aligned expert ranges before trimming with synchronized expert counts.

Quantization

• Guard cast-back and lossless/per-token scaling-factor loads and stores against the real scaling-factor grid dimensions.
• Guard packed E5M6 output stores for token and hidden tail tiles.
• Use FP32 fragments for per-token absmax reductions before computing reciprocal scales, instead of using in_config.dtype as the accumulator dtype.

Test and benchmark tooling

• Add run_cutedsl_tests.sh for one-command CuTeDSL correctness, benchmark, or combined full runs.
• Preserve externally provided CUDA_VISIBLE_DEVICES under xdist instead of always overriding it.
• Add optional per-worker TileLang cache directories and failure/trace report files for long full-suite runs.
• Add an opt-in missing-baseline override for benchmark sweeps that generate new JSONL data before a baseline exists.
• Make the benchmark timer backend configurable through TK_BENCHMARK_BACKEND.

Why these are TileKernels fixes

These changes address TileKernels assumptions about launch geometry, tensor bounds, allocation sizing, and reduction accumulator dtype. They are not CuTeDSL-version-specific workarounds:

• The Engram issue follows from the kernel assuming a persistent block count divisible by four while the project requirements state SM90 support.
• The MoE and quant changes guard logical tensor bounds when the implementation intentionally uses aligned loop extents.
• The FP32 absmax change fixes the accumulator dtype used to derive quantization scales, independent of backend code generation details.

Validation

Validated on NVIDIA H100 PCIe with the local TileLang CuTeDSL backend after the corresponding TileLang CuTeDSL fixes:

• Full correctness coverage: passed after rerunning the large OOM-sensitive cases in a single-worker H100 run.
• Full benchmark coverage: passed for run-through status with benchmark JSONL emission enabled.
• Current PR sanity checks:
  • git diff --check HEAD~4..HEAD
  • .venv-tk-test/bin/python -m ruff check ... on changed Python files
  • .venv-tk-test/bin/python -m py_compile ... on changed Python files
  • bash -n run_cutedsl_tests.sh

Commit Structure

• [BugFix][Engram] Support non-divisible SM90 block counts
• [BugFix][MoE] Guard fused mapping edge tiles
• [BugFix][Quant] Guard scale tails and FP32 absmax
• [Test] Add reusable CuTeDSL full-run controls

[JayceSu98]

Superseded by the three split bugfix PRs: #2 for Engram, #3 for MoE, and #4 for Quant. The previous combined PR also contained reusable test tooling, which is intentionally left out of the split bugfix review.