[DSv4.1] Fuse native-layout vision RoPE backward and QKV gradient packing

[YangXu1990uiuc]

The published DSv4.1 vision attention path produces BF16 Q/V gradients in HTD storage and K gradients in HDT storage. This adds VisionRoPEBackward and its allocation wrapper with backend="frost": one SM100 kernel reads those layouts directly, applies split-half rotary backward, and writes contiguous packed dQKV. It removes the separate K-layout copy and fuses Q/K rotation with V packing.

The API supports dynamic token counts without recompilation, caller-owned output, explicit streams, and CUDA Graph capture. It validates dtype, layout, alignment, device, and output overlap. Trigonometric tables use FP32; multiply and add/subtract round separately before BF16 output. The implementation targets 16 heads and head dimension 64.

The mathematical contract follows DeepSeek-V4.1-Flash vision.py. The GPU addressing, shared-memory transpose, and fused packing are authored for this implementation.

Validation:

• 29 focused GPU tests pass on SM100, covering six model-derived image geometries, tail tiles, fresh buffers, dynamic T, invalid layouts/dtypes/alignment, overlap, streams, and Graph replay.
• B200 validation loads the actual checkpoint's first vision block. All 24 ordinary/checkpointed forward-and-backward comparisons pass across six geometries and two input banks; intermediate operands and the public FE kernel route are verified.
• A separate B200 run validates the complete 32-layer vision tower and aligner with actual checkpoint weights, six model-derived geometries, and two input banks per geometry. Original 3D math attention and RoPE forward remain unchanged; non-reentrant checkpointing is enabled per block. Each run checks the output, all 32 hidden states, input gradient, and all 263 parameter gradients against the source implementation. All 297 tensors match exactly numerically; outputs and hidden states also match bitwise. Source-repeat and Torch-formula adapter controls pass, and profiling confirms one native FROST backward call per layer (384 calls across 12 cases). The compiled PTX and cubin are identical to the operator benchmark.
• Two separate B200 runs qualify the kernel. The public API run compares original autograd, Transformer Engine, FlashInfer, four compiled Torch variants, and a fused FROST control that first copies dK. All required layout conversions and packed-output writes are timed. Each cache regime uses eight paired blocks with five samples per provider and block.

The following backward operator speedups use the fastest same-run control, which was the copy-dK-plus-FROST implementation:

Vision tokens	Graph GPU speedup (hot / evicted)
1521	1.36x / 1.37x
1610	1.36x / 1.34x
4070	1.39x / 1.38x
5476	1.40x / 1.40x
8418	1.41x / 1.42x
8649	1.39x / 1.39x

Across these shapes, public-API Graph GPU speedup is 1.34–1.42x and Graph wall-time speedup is 1.29–1.42x. Every paired block exceeds 1.05x for both Graph clocks and cache regimes. Relative to the fastest external provider, Graph GPU speedup is 1.76–1.91x. Eager wall-time medians improve 1.03–1.13x, with smaller and noisier gains.

A separate B200 timing run covers forward plus checkpointed backward of the full 32-layer vision tower and aligner, using the same actual checkpoint, six image geometries, two input banks, and all nine providers. Each geometry uses eight paired blocks with three samples per provider; compilation is excluded and outputs/gradients are rechecked after timing. Relative to the fastest same-run control, median whole-phase wall-time ratios are 0.999–1.003x, and CUDA-event span ratios are 0.999–1.003x; against only the numerically exact controls, wall-time ratios are 1.001–1.006x. None of the six geometries shows a stable >5% whole-phase gain. These results do not demonstrate an overall vision-training speedup.

The performance benefit demonstrated here is confined to the RoPE backward-and-pack operator. The vision tests retain the source 3D math attention path and exclude optimizer updates and the language backbone. Only the kernel, public API, user documentation, and focused tests are included.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}