[pretuned] Optimize fused_qk_norm_rope on B200

[yushangdi]

Summary

Optimize the pretuned fused_qk_norm_rope kernel for full rotary dimensions and synchronize its B200 heuristic with vLLM PR vllm-project/vllm#50466.

For rotary_dim == head_dim, the kernel now:

• loads the two rotary halves directly from QKV;
• computes the full-head RMS reduction from those register-resident values;
• applies Q/K weights and RoPE without an intermediate global-memory round trip;
• stores each output half once.

The emitted full-rotary Triton contains 2 QKV loads, 2 QKV stores, and no tl.debug_barrier().

Partial rotary dimensions retain the existing full-head normalization path because RMSNorm must also update the non-rotary tail.

Config changes

• Import the measured B200 schedules from the corresponding vLLM change.
• Use one-head/one-warp schedules for 1-2 tokens and two-head/one-warp schedules through 256 tokens.
• Retain the existing persistent schedules where they remain faster at large token counts.
• Replace stale per-operation indexing lists with scalar pointer indexing after the memory-operation graph changed.
• Add the measured SM100 performance expectation to the pretuned performance gate.

Related work and duplicate check

This is not a duplicate of #3081. That PR fixed the general correctness hazard by inserting a barrier for intra-loop store-to-load dependencies. This PR removes that dependency from the common full-rotary kernel path, eliminating both the QKV store/reload and the barrier while preserving #3081 as the compiler-level safety mechanism for kernels that still require such dependencies.

Searches for fused qk norm rope, store forwarding, and redundant store load found no other open Helion PR implementing this dataflow or these B200 schedules.

Tests

CUDA_VISIBLE_DEVICES=7 .venv/bin/python -m pytest \
  test/test_pretuned_kernels.py::TestPretunedKernelsCorrectness::test_fused_qk_norm_rope -q

CUDA_VISIBLE_DEVICES=7 .venv/bin/python -m pytest \
  test/test_pretuned_kernels.py::TestPretunedKernelsPerformance::test_fused_qk_norm_rope -q -s

uv --offline tool run --from ruff ruff check \
  pretuned_kernels/fused_qk_norm_rope/fused_qk_norm_rope.py \
  pretuned_kernels/fused_qk_norm_rope/_helion_aot_fused_qk_norm_rope_cuda_sm90.py \
  pretuned_kernels/fused_qk_norm_rope/_helion_aot_fused_qk_norm_rope_cuda_sm100.py \
  test/test_pretuned_kernels.py

Results:

• Correctness gate: 1 passed
• Performance gate: 1 passed
• B200 sweep versus vLLM CUDA: 21/21 wins, 1.3021x geomean, 1.6855x best speedup
• Ruff check and format check: passed
• git diff --check: passed

./lint.sh could not run because Ruff/Pyrefly are not installed in this checkout, and repository instructions prohibit networked installs. Offline cached Ruff passed; Pyrefly remains an unrun test gap. H100 structural schedules are retained, but H100 performance was not rerun after changing the memory-operation indexing policy.

AI assistance

AI assistance was used to inspect generated Triton, develop and validate the kernel dataflow, tune B200 schedules, and draft this description. The human submitter must review the changes and results before requesting review.