[garrett361]

Adds a cudnn_flashmla option to model.dsa_backend for DeepSeek V4 that runs the same FlashMLA sparse prefill forward vLLM uses for DeepSeek V4 prefill, so the trainer and inference compute attention with the same kernel. The backward is the cuDNN frontend's DSA kernel; it runs on SM90, SM100 and SM103, and tilelang stays the default.

• For DeepSeek V4, whose slot width changes with sequence length and packing, cudnn_dsa_patch stops cuDNN
  recompiling its backward for each new width. GLM-5's backward uses one fixed width and is unaffected.
• On B300 (SM103), both DSv4 test files pass and the kernels match a float32 reference; no end-to-end run there yet.
• On H200 at the B300 changes (6a1cae560), both DSv4 test files pass again, and a 3-step end-to-end run (cp 8,
  ep 8) compiles 3 cuDNN kernels per rank in step 1 and matches the earlier losses, step time and peak memory.
• Fixes GLM-5's SM100 cudnn_flashmla backward, which omitted cuDNN's required workspace and always raised.
• tilelang stays the shared dsa_backend default and configs/advanced/deepseek-v4-flash/sft.toml opts in to
  cudnn_flashmla, following perf(dsa): cuDNN + FlashMLA kernels for GLM-5 sparse attention (model.dsa_backend) #3829, which did the same for the GLM-5.3 configs. FlashMLA only ships an x86_64 wheel
  in an optional extra, so a cudnn_flashmla default would break aarch64 and installs without the extra. We can flip
  the default instead if reviewers prefer.

DeepSeek-V4-Flash first 6 layers, real weights, cp 8, ep 8, one 8xH200 node. Sparse attention GPU time per step, rank 0 profiler trace at 8k tokens/GPU:

kernels	tilelang	cudnn_flashmla	speedup
forward, incl. recompute (48 calls)	56.3 ms	26.5 ms	2.12x
backward (24 calls)	71.3 ms	36.7 ms	1.94x
all sparse attention	136.0 ms	76.7 ms	1.77x

Sparse attention is under 5% of the step here, so the end-to-end gain is smaller. Median steady-state step time:

tokens/GPU	tilelang	cudnn_flashmla	speedup	max loss diff vs tilelang	max loss diff, two tilelang runs
8k	2.86 s	2.80 s	1.02x	1.0e-3	1.7e-3
16k	4.59 s	4.36 s	1.05x	7.3e-4	1.4e-3

Peak memory is equal, and each rank compiles 3 cuDNN kernels, all in step 1. Baseline main; results and config at 524ec668e.

🤖 Generated with Claude Code