Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
Signed-off-by: Johnny <johnnynuca14@gmail.com>
Signed-off-by: neweyes <328719365@qq.com>

…es on SM12x (#2913)

### Summary

- Add missing `-DCUTLASS_ENABLE_GDC_FOR_SM100=1` compile flag to all
CUTLASS fused MoE JIT modules (SM100/SM103/SM120) and
`-DCUTLASS_ENABLE_GDC_FOR_SM90=1` to SM90 modules
- Sync nv_internal `grid_dependency_control.h` with upstream CUTLASS to
support SM100/SM103/SM110/SM120/SM121 GDC
- Add `-DCUTLASS_ENABLE_GDC_FOR_SM90=1` to FP8 blockscale GEMM SM90
module

### Problem

Random `cudaErrorIllegalInstruction` crashes on DGX Spark (SM121) and
RTX 50-series (SM120) when running NVFP4 MoE models (e.g., Nemotron,
Qwen3.5-122B) under load. The crashes are intermittent and worsen with
longer context lengths and higher concurrency.

**Root cause:** PR #2780 fixed the missing GDC compile flags for GEMM
modules (`flashinfer/jit/gemm/core.py`), but the **CUTLASS fused MoE
modules** in `flashinfer/jit/fused_moe.py` and the **FP8 blockscale GEMM
module** were not fixed. This is the exact same class of bug as #2708.

Without `-DCUTLASS_ENABLE_GDC_FOR_SM100=1`, CUTLASS's
`grid_dependency_control.h` compiles `wait_on_dependent_grids()` and
`launch_dependent_grids()` as **empty no-ops**:

```cpp
CUTLASS_DEVICE void wait_on_dependent_grids() {
#if (defined(CUTLASS_GDC_ENABLED))   // ← not defined without the flag
  asm volatile("griddepcontrol.wait;");
#endif
}
```

Meanwhile, the host-side code still sets
`programmaticStreamSerializationAllowed = true` (PDL enabled) via
`device_support_pdl()` which returns `True` for all `major >= 9`,
including SM12x. This means:

1. **Host enables PDL** → CUDA runtime overlaps consecutive kernels
2. **Device GDC barriers are no-ops** → No synchronization between
overlapping kernels
3. **Race condition** → Dependent kernel reads stale global memory →
corruption → `cudaErrorIllegalInstruction`

The crash is random because it depends on exact kernel scheduling
timing, which varies per request.

### Fix

**`flashinfer/jit/fused_moe.py`** — Added GDC flags to all CUTLASS fused
MoE modules:

| Module | Flag | Architectures Covered |
|---|---|---|
| `gen_cutlass_fused_moe_sm120_module()` |
`-DCUTLASS_ENABLE_GDC_FOR_SM100=1` | SM120, SM121 |
| `gen_cutlass_fused_moe_sm103_module()` |
`-DCUTLASS_ENABLE_GDC_FOR_SM100=1` | SM103, SM120, SM121 |
| `gen_cutlass_fused_moe_sm100_module()` |
`-DCUTLASS_ENABLE_GDC_FOR_SM100=1` | SM100, SM110, SM120, SM121 |
| `gen_cutlass_fused_moe_sm90_module()` |
`-DCUTLASS_ENABLE_GDC_FOR_SM90=1` | SM90 |
| `gen_trtllm_gen_fused_moe_sm100_module()` |
`-DCUTLASS_ENABLE_GDC_FOR_SM100=1` | SM100+, SM120, SM121 |

**`flashinfer/jit/gemm/fp8_blockscale.py`** — Added
`-DCUTLASS_ENABLE_GDC_FOR_SM90=1` to
`gen_fp8_blockscale_gemm_sm90_module()`.

**`csrc/nv_internal/.../grid_dependency_control.h`** — Synced with
upstream CUTLASS
(`3rdparty/cutlass/include/cutlass/arch/grid_dependency_control.h`) to
add SM100+ GDC support. Previously only handled SM90, so any nv_internal
TensorRT-LLM code compiled for SM12x would have GDC barriers silently
compiled as no-ops.

### Why `-DCUTLASS_ENABLE_GDC_FOR_SM100=1` covers SM12x

CUTLASS uses a single flag for the entire Blackwell family. From
`grid_dependency_control.h`:

```cpp
#if(CUDA_BARRIER_ENABLED && defined(CUTLASS_ENABLE_GDC_FOR_SM100) && defined(__CUDA_ARCH__) && \
    ((__CUDA_ARCH__ == 1000 && ...) ||   // SM100
     (__CUDA_ARCH__ == 1030 && ...) ||   // SM103
     (__CUDA_ARCH__ == 1100 && ...) ||   // SM110
     (__CUDA_ARCH__ == 1200 && ...) ||   // SM120 (RTX 50-series)
     (__CUDA_ARCH__ == 1210 && ...)))    // SM121 (DGX Spark)
#define CUTLASS_GDC_ENABLED
```

### Why SM90 GDC flag was NOT added to SM100+ modules

PR #2716 attempted to add both `-DCUTLASS_ENABLE_GDC_FOR_SM90=1` and
`-DCUTLASS_ENABLE_GDC_FOR_SM100=1` to all modules. It broke AOT builds
because `sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp`
checks `CUTLASS_ENABLE_GDC_FOR_SM90` and calls
`scheduler.is_last_tile()` — a method not present on the SM120
scheduler. PR #2780 corrected this by using only the SM100 flag for
SM100+ modules. This PR follows the same approach.

### Related

- #2708 — Original issue: missing GDC flags cause PDL race condition
- #2716 — First fix attempt (reverted — broke AOT)
- #2780 — Corrected fix for GEMM modules only
-
[vllm-project/vllm#38423](vllm-project/vllm#38423)
— NVFP4 bugfix on DGX Spark
- [NVIDIA/cutlass#3121](NVIDIA/cutlass#3121) —
K=64 block-scaled GEMM tiles (separate issue)

### Test plan

- [x] Clear JIT cache: `rm -rf ~/.cache/flashinfer/`
- [x] Run NVFP4 MoE model on SM121 (DGX Spark) with 128K context under
load — verify no `cudaErrorIllegalInstruction`
- [x] Run NVFP4 MoE model on SM120 (RTX 50-series) with concurrent
requests — verify no NaN/garbage output
- [x] Verify `CUDA_LAUNCH_BLOCKING=1` workaround is no longer needed
- [x] AOT build with `FLASHINFER_CUDA_ARCH_LIST="12.1a"` completes
without errors
- [x] SM90 (Hopper) fused MoE tests pass: `pytest tests/moe/`
- [x] SM100 GEMM tests still pass (no regression from existing GDC
flags)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Expanded GPU kernel compilation support: enabled additional
optimizations for NVIDIA SM100 and SM90 GPUs, activating
dependency-control optimizations where available.
* Updated JIT/GEMM build configs to include these architecture-specific
compile options, improving performance and compatibility on supported
hardware.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
Signed-off-by: Johnny <johnnynuca14@gmail.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>

Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
Signed-off-by: Johnny <johnnynuca14@gmail.com>