diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index 660e29f22..f8c5ebba2 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -112,12 +112,14 @@ jobs: - "python/sglang/kernels/ops/diffusion/**" - "test/registered/kernels/ops/diffusion/**" - "test/registered/kernels/benchmark/diffusion/**" + - "test/registered/kernel/diffusion/**" - "python/sglang/cli/**" jit_kernel: - ".github/workflows/pr-test.yml" - ".github/workflows/pr-test-jit-kernel.yml" - "python/pyproject.toml" - "test/registered/kernels/**" + - "test/registered/kernel/diffusion/**" # sglang.kernels is the migrated kernel namespace (RFC #29630 / #30044); the # base-b-kernel suites import it directly, so kernel edits must run them. - "python/sglang/kernels/!(*.md)" diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index e159e6ddd..c49f66eb5 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -800,7 +800,84 @@ SGLang port. Use base MiniMax-H3 when output quality matters more than latency. -## 7. Sampling and output controls +## 7. VDN-H3: hybrid attention, 8-step distill + +[OpenVDN/vdn-minimax-h3](https://huggingface.co/OpenVDN/vdn-minimax-h3) +(Video DeltaNet MiniMax-H3, [openvdn.github.io](https://openvdn.github.io/)) +replaces every DiT block's dense self-attention with a hybrid of two branches: +an exact, gated softmax over a chunk-aligned frame window (chunk 5, radius 1, +the first and last frames dense as anchors, text and audio dense both ways) +and a frame-wise linear-attention branch (the Video Delta rule, scanned +forward and backward over frames) that covers exactly the rest. The released +`stage-dmd-step-250` checkpoint is an 8-NFE DMD2 distill and adds a 4.3 GB +linear branch plus two small LoRA adapters on top of the untouched MiniMax-H3 +backbone. The weights are the FL2VA partition, so the deployment serves `t2va` +and `fl2va` (first / last keyframes); the keyframe rows are attended densely, +like text and audio. `ref2va` was not trained and is rejected. The weights +inherit the MiniMax-H3 Community License, including its territorial exclusions. + +Pass the repository directly to `--model-path`. A registered model overlay +materializes the base-H3 layout once: both LoRA adapters are prefused into the +transformer weights (a real 62 GB write), the linear branch is attached as an +extra shard, the Qwen3-VL conditioner is linked from `MiniMaxAI/MiniMax-H3`, +and the video VAE is re-serialized as for FastH3. Point +`SGLANG_DIFFUSION_CACHE_ROOT` at a volume with at least 90 GB free before the +first launch. + +```bash 4×B200 hybrid window attention +sglang serve \ + --model-path OpenVDN/vdn-minimax-h3 \ + --num-gpus 4 \ + --attention-backend hybrid_window_attn_h3 \ + --performance-mode speed \ + --warmup-num-frames 345 \ + --warmup-resolutions 1344x768 \ + --port 30010 +``` + +`--warmup-num-frames` / `--warmup-resolutions` make the startup warmup run +at the clip length and canvas you will serve (here the 14.375 s paper +workload). Without them H3 warms up on a 5-second clip, and the first +forward of the first longer request pays allocator growth and kernel setup +in every DiT block (2 to 3 s on this workload, for base H3 as well). + +Requests use the same asynchronous video endpoint as the base model, with +`task: "t2va"`, `conditions: []` (or `task: "fl2va"` with the keyframe +conditions of the [FL2VA request](#4-generate-video-and-audio)), and a target such as +`{"short_edge": 768, "aspect_ratio": "16:9", "duration_seconds": 14.375}`. Each +keyframe adds about 2,000 global rows (the conditioner's image tokens plus the +latent keyframe), dense both ways in the window softmax, so an `fl2va` forward +costs more than a `t2va` one. The +request default is `num_inference_steps: 9`: nine points on the standard +shift-12/shift-3 sigma grid, i.e. the eight distilled DiT evaluations (VDN +counts NFEs, SGLang counts grid points). Any other step count is rejected. The +paper workload is 1344×768 at 24 fps for 14.375 s: 345 frames, already 17n+5 +aligned, 102 latent frames and about 104.5k packed rows. + +`hybrid_window_attn_h3` is required for the transformer: a dense backend on +these weights would silently skip the linear branch and the softmax gates and +produce the wrong model, so it is rejected. The window softmax runs as a union +of dense FlashAttention varlen calls, exact to bf16 rounding. The +linear branch runs on fused Triton kernels (temporal conv + SiLU + L2 norm, +statistics prologue, gated RMSNorm epilogue). On Blackwell (SM100 / SM103, +and SM120 such as the RTX PRO 6000) the transformer defaults to online MXFP8 +(`--quantization fp8` selects it too, `--quantization bf16` opts out); the +gates, beta, alpha and the conv stay bf16. Before SM100 `--quantization fp8` +is the per-channel fp8 path of base H3 (SM89 and SM90); SM80 has no fp8 +tensor cores, so it runs the bf16 DiT (62 GB, budget for +`--layerwise-offload-components` or `--dit-layerwise-offload`). Ampere and +Ada run the window on FA3's Sm80 mainloop, which is FA2-class throughput; they +are enabled but not benchmarked. Ulysses sequence parallelism is supported; +`--model-variant`, `quality: "high"`, `--ring-degree` greater than 1, +`torch.compile`, and breakable CUDA graph execution are rejected. See +[Attention Backends](/docs/sglang-diffusion/attention_backends) for the +backend options. + +Measured latencies for the 4× B200 recipe are in +[VDN-H3 on B200](#vdn-h3-on-b200); single-card and PCIe multi-card numbers +for the RTX PRO 6000 are in [VDN-H3 on RTX PRO 6000](#vdn-h3-on-rtx-pro-6000). + +## 8. Sampling and output controls MiniMax-H3 supports more than one output per prompt. The video API accepts `num_outputs_per_prompt` (or OpenAI-compatible `n`) from 1 through 10. Offline @@ -975,7 +1052,7 @@ has completed, but the `quality: "high"` path above remains fail-closed to the audited 4×H200 workload. -## 8. Feature contracts and advanced recipes +## 9. Feature contracts and advanced recipes The generated command already contains the recommended topology and encoder setting. Use the detailed reference below only when applying an optional @@ -1269,7 +1346,7 @@ fold decision is not node-boundary aware: -## 9. Configuration notes +## 10. Configuration notes - MiniMax-H3 produces the canonical 24 fps output; request duration is expressed through `target.duration_seconds`. - `target.duration_seconds` must be between 4 and 15 seconds, inclusive. The command picker defaults to the verified 5-second profile. @@ -1288,7 +1365,7 @@ fold decision is not node-boundary aware: - `speed` keeps model components resident, while `auto` applies the model-aware 120 GiB residency threshold. `memory` prioritizes avoiding OOM and includes the executable VAE decoder in its default layerwise set. A measured recipe with sufficient headroom can opt into `--component-residency vae=resident`; the 2×H100 CI recipe does this because the VAE's 4.8 GiB/GPU cost avoids repeated decoder transfers during tiled decode. DiT residency and prefetch knobs remain scoped to the DiT. Use `speed` only after confirming that the complete target workload fits. - Breakable CUDA graph execution is an explicit opt-in, not part of the recommended `speed` preset. It requires `--enable-breakable-cuda-graph`, every served size in `--warmup-resolutions`, and `--bcg-text-buckets` that cover the live H3 condition sequence. The validated 1344×768 Ref2VA recipe uses 5504; other task profiles and reference sets may need a different value. It preserves eager output for matching captured signatures, but graph capture consumes additional GPU memory and may provide little latency benefit when Ulysses attention and collectives dominate, so benchmark it on the target topology before enabling it. -## 10. Benchmarks +## 11. Benchmarks The picker exposes resident and FSDP profiles on NVIDIA datacenter GPUs. GPU counts are properties of the selected recipes, not a claim that every platform @@ -1452,6 +1529,107 @@ TP2 + Ulysses2 (3.42 s, 62,290 MB), FSDP + Ulysses4 (3.42 s, 50,984 MB), and online `--quantization fp8` (2.93 s, 64,204 MB) trade a little latency for peak memory. +### VDN-H3 on B200 + +A 4× B200 (SM100, 183 GB) host served [VDN-H3](#7-vdn-h3-hybrid-attention-8-step-distill) +on the paper workload: 1344×768 at 24 fps with audio for 14.375 s (345 frames, +102 latent frames, about 104k packed rows), `task: "t2va"`, +`num_inference_steps: 9` (8 DiT forwards), seed 1000, `hybrid_window_attn_h3` +with the decomposed window kernel, eager, and the warmup run at the served +clip shape (`--warmup-num-frames 345 --warmup-resolutions 1344x768`), so +every forward of the served request is steady state; without those flags the +first forward pays about 3 s of allocator growth and kernel setup. "Steady +s/NFE" is the mean of forwards 2 to 8. The OpenVDN reference rows ran the released +inference stack (`8nfe_tuned_fp8.yaml`, `infer_ulysses.py`) on the same host +and the same clip length: + +| Config | Steady s/NFE | Denoise (8 forwards) | Decode | Peak/GPU | +| --- | ---: | ---: | ---: | ---: | +| OpenVDN published, FP8, 8× B200 Ulysses 5+3 ([openvdn.github.io](https://openvdn.github.io/)) | 1.40 | 11.2 s | – | – | +| OpenVDN reference, BF16 (`8nfe_tuned.yaml`), 1× B200 | 7.91 | 63.3 s | – | – | +| OpenVDN reference, FP8, 1× B200 | 6.47 | 51.7 s | – | – | +| OpenVDN reference, FP8, 4× B200 Ulysses (standard / 3+1 branch-parallel) | 2.68 / 2.61 | 21.4 / 20.9 s | – | – | +| SGLang, BF16, 1× B200, DiT layerwise offload (auto policy) | 7.86 | 64.1 s | 9.5 s | 147,608 MB | +| SGLang, FP8, 1× B200, DiT layerwise offload (auto policy) | 7.49 | 61.7 s | 9.4 s | 111,616 MB | +| SGLang, FP8, 1× B200, DiT resident (`--layerwise-offload-components text_encoder`) | 7.49 | 61.1 s | 9.3 s | 63,922 MB | +| SGLang, `--quantization bf16`, 4× B200 Ulysses4, `--performance-mode speed`, served-shape warmup | **2.53** | **20.3 s** | 3.5 s | 97,894 MB | +| SGLang, FP8, 4× B200 Ulysses4, `--performance-mode speed`, served-shape warmup | **2.34** | **18.7 s** | 3.4 s | 63,022 MB | +| SGLang default (online `mxfp8`), 8× B200 Ulysses8, `--performance-mode speed`, served-shape warmup | **0.93** | **7.45 s** | 1.5 s | 77,704 MB | +| SGLang, per-channel fp8 weight scales (the online `fp8` path of other models), 8× B200 Ulysses8 | 1.05 | 8.4 s | 1.5 s | 76,828 MB | + +On 8× B200 the SGLang Ulysses8 path runs the paper workload at 0.93 s/NFE +(7.5 GPU-seconds per NFE) against the published 1.40 s/NFE of OpenVDN's 5+3 +branch-parallel layout; the whole request (text encoding, 8 forwards, joint +decode) completes in 9.3 s after warmup. At this GPU count the step is +launch- and copy-bound in the linear branch rather than FLOP-bound: the scans +fold the 102 frames into per-chunk composites plus one chain over the chunks +with both directions per launch (the chunked window only reads states at +chunk boundaries), q/k/v and the per-head scalars travel as four async +field-major all-to-alls that land contiguous (no relayout copies for the +window K/V gathers or the branch's conv), only the 128-wide output-gate +hidden crosses the fabric (the gate's `up` runs on the head shard), the frame +mean is a reshape-sum instead of an atomic index_add, and the transformer +runs online `mxfp8` by default (e4m3 with one E8M0 scale per 32 elements into +cuBLASLt's block-scaled GEMM; the activation quant is fused into the adaLN +modulation and SwiGLU producers, so no standalone quant pass runs). The +prompt's text state joins +the frames' Cholesky batch as a virtual frame. OpenVDN's branch-parallel layout was measured slower here (1.39 +s/NFE at 5+3, 1.49 at 6+2): with 11 to 12 heads per softmax rank the window +attention and its gathers grow faster than the linear ranks shrink; splitting +each rank's heads into two pipelined all-to-all groups also lost (1.10 s/NFE) +to smaller attention kernels and doubled branch launches. The profile of one +rank is 31% window FlashAttention, 24% fp8 GEMM, 12% NCCL, the rest small +kernels. +Against the published 8× B200 headline (1.40 s/NFE, 11.2 +GPU-seconds per NFE), the SGLang 4× B200 FP8 run spends 9.8 GPU-seconds per +NFE, with half the GPUs and half the all-to-all fan-out. On the same host and GPU count the +SGLang path is 10-13% faster per NFE than the released stack (2.34 vs 2.61 / +2.68 s), and the 8-forward denoise is 18.7 s against their 20.9 / 21.4 s. In BF16 the single-GPU path matches the released tuned BF16 stack +(7.86 vs 7.91 s/NFE); the single-GPU FP8 gap (7.49 vs 6.47 s) is the online +FP8 GEMM path, not the hybrid attention. + +On one B200 the auto memory policy streams the DiT layers from host memory +(the 62 GB DiT, the 66 GB conditioner and the VAEs do not all fit resident +with the 120 GB headroom the policy keeps); keeping the FP8 DiT resident with +`--layerwise-offload-components text_encoder` changes nothing measurable, so +the single-GPU rows are compute-bound. The BF16 DiT does not fit resident at +this clip length on one 183 GB card (the 104k-row activations alone take +about 100 GB). Per DiT block the hybrid attention costs about 117 ms at this +shape (window FlashAttention 46 ms, projections 24 ms, the linear branch and +its fused kernels the rest). A static-tile block-sparse Triton kernel was +measured at 226 ms against 169 ms for the decomposed FlashAttention path and +is not shipped. + +### VDN-H3 on RTX PRO 6000 + +The same paper workload (1344×768, 345 frames, 8 DiT forwards, seed 1000, +`hybrid_window_attn_h3`, `--performance-mode speed`, served-shape warmup) on +RTX PRO 6000 Blackwell Server Edition cards (SM120, 96 GB, PCIe, no NVLink). +The window runs on FA4's sm120 kernel and the transformer defaults to online +`mxfp8` (a cutlass sm120 block-scaled GEMM); `--quantization fp8` maps to it. +One card needs `--layerwise-offload-components text_encoder`: the fp8 DiT +(33 GB) and the bf16 text encoder (48 GB) do not both fit, and the whole-module +`--text-encoder-cpu-offload` re-homes the encoder in one piece and runs out of +memory when it is used. Multi-card Ulysses goes over PCIe; the all-to-all is +the bulk of the step there (54% of GPU time at 8 cards), so the per-card +efficiency drops with count. A 2-card run should use a pair on one PCIe +switch (`nvidia-smi topo -m` shows `PIX`): a pair across the CPU root complex +measured 15.7 to 17.0 s/NFE on a shared host. + +| Config | Steady s/NFE | Denoise (8 forwards) | Decode | Peak/GPU | +| --- | ---: | ---: | ---: | ---: | +| 1× RTX PRO 6000, `--layerwise-offload-components text_encoder` (default `mxfp8`) | **20.27** | **160.0 s** | 20.6 s | 95,790 MB | +| 1× RTX PRO 6000, same, per-channel `fp8` path (before this mapping) | 22.85 | 180.3 s | 20.7 s | 95,160 MB | +| 2× RTX PRO 6000 Ulysses2, one PCIe switch (default `mxfp8`) | 13.61 | 107.4 s | 10.6 s | 93,846 MB | +| 2× RTX PRO 6000 Ulysses2, one PCIe switch, per-channel `fp8` | 14.97 | 118.1 s | 10.6 s | 94,700 MB | +| 4× RTX PRO 6000 Ulysses4 (default `mxfp8`) | 7.98 | 63.0 s | 5.6 s | 82,178 MB | +| 4× RTX PRO 6000 Ulysses4, per-channel `fp8` | 8.23 | 65.1 s | 5.6 s | 82,422 MB | +| 8× RTX PRO 6000 Ulysses8 (default `mxfp8`) | 5.43 | 43.1 s | 3.5 s | 77,704 MB | +| 8× RTX PRO 6000 Ulysses8, per-channel `fp8` | 5.47 | 43.4 s | 3.5 s | 77,236 MB | + +For comparison the same code on B200 measures 6.23 / 3.34 / 1.73 / 0.89 s/NFE +at 1 / 2 / 4 / 8 cards (NVLink, 87% parallel efficiency at 8). + ### H200 topology comparison The same four-card H200 host completed both lossless resident placements with diff --git a/docs/docs/sglang-diffusion/attention_backends.mdx b/docs/docs/sglang-diffusion/attention_backends.mdx index e9536aaee..5e138113a 100644 --- a/docs/docs/sglang-diffusion/attention_backends.mdx +++ b/docs/docs/sglang-diffusion/attention_backends.mdx @@ -89,6 +89,11 @@ For SGLang-native pipelines, the CLI accepts the lowercase names of `AttentionBa `VIDEO_SPARSE_ATTN_H3` Video Sparse Attention for MiniMax-H3 / FastH3 (VSA-H3). In-tree Triton block-sparse kernel (SM90 / SM100 / SM103); no external package. Configure via --attention-backend-config. + + `hybrid_window_attn_h3` + `HYBRID_WINDOW_ATTN_H3` + VDN-H3 hybrid attention: chunk-aligned window softmax (exact, gated) plus the Video Delta linear branch over the window complement. The window runs as a union of FlashAttention varlen calls (FA4 on SM100 / SM103 / SM120, FA3 on SM90; on SM80 / SM86 / SM89 the same FA3 build runs its Sm80 mainloop at FA2-class speed). Configure via --attention-backend-config. + `vmoba_attn` `VMOBA_ATTN` @@ -309,6 +314,62 @@ VSA-H3 constraints: - Ulysses sequence parallelism is supported; `--ring-degree` greater than 1, `torch.compile`, and breakable CUDA graph execution are rejected. +**Hybrid window attention for VDN-H3 (`hybrid_window_attn_h3`)** + +[VDN-H3](/cookbook/diffusion/MiniMax/MiniMax-H3#7-vdn-h3-hybrid-attention-8-step-distill) +replaces every DiT block's dense self-attention with two branches. The +softmax branch is an exact softmax over a chunk-aligned frame window (frame +`t` belongs to chunk `t // 5` and attends to chunks `c - 1 .. c + 1`; frames +0 and F-1 are dense as rows and columns; every text / audio pair stays dense) +scaled by a per-head sigmoid gate. The linear branch (a frame-wise Video Delta +rule recurrence, forward and reverse over frames) covers exactly the window's +complement and is driven by the attention module. The mask is request-static, +so the metadata is built once per request. The checkpoint's +`transformer/config.json` carries the window geometry; the backend reads it. + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ParameterTypeDescriptionDefault
`vdn_h3_dense_smoke``bool`Allows a dense transformer backend (`fa`) on the VDN-H3 weights for the base-H3 + LoRA equivalence smoke. The gates and the linear branch are skipped: a study, not a sample.`false`
`vdn_max_gather_rows``int`Upper bound on the gathered K/V rows per window FlashAttention call; consecutive chunk groups fill one call up to it. Splitting changes no query's kept key set, only the pass count and the gather's peak memory.`200000`
+ +Hybrid window attention constraints: + +- VDN-H3 requires this backend for the transformer; a dense backend on these + weights would silently skip the linear branch and the gates, so it is + rejected unless `vdn_h3_dense_smoke` is set. The token refiner, text + encoder, and VAEs keep dense attention. +- Base MiniMax-H3 and FastH3 checkpoints have no linear branch and are + rejected by this backend. +- Ulysses sequence parallelism is supported (QK-norm + RoPE run after the + all-to-all on the head shard); `--ring-degree` greater than 1, + `torch.compile`, and breakable CUDA graph execution are rejected. + **V-MoBA (`vmoba_attn`)** diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 4ada585a9..ec372d063 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -373,7 +373,7 @@ For a pipeline whose primary DiT is named `transformer`, the shorter MiniMax-H3 supports this path while preserving its required FP32 patch, timestep, and output projections. See the -[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#8-feature-contracts-and-advanced-recipes) +[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#9-feature-contracts-and-advanced-recipes) for its distributed serving recipe. ### MXFP4 Online Quantization @@ -432,7 +432,7 @@ projections take that path. `kitchen_int8` is approximate and is not a consistency ground-truth mode. The BF16 path is unchanged when `comfy-kitchen` is not installed. See the -[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#8-feature-contracts-and-advanced-recipes) +[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#9-feature-contracts-and-advanced-recipes) for the 24 GB offload recipe, including why `vae` must stay out of `--layerwise-offload-components`. diff --git a/docs/src/snippets/diffusion/model-catalog.jsx b/docs/src/snippets/diffusion/model-catalog.jsx index 25c95d428..5779f6cfd 100644 --- a/docs/src/snippets/diffusion/model-catalog.jsx +++ b/docs/src/snippets/diffusion/model-catalog.jsx @@ -190,6 +190,11 @@ export const DiffusionModelCatalog = ({ category }) => { ], cookbook: "/cookbook/diffusion/MiniMax/MiniMax-H3#6-fasth3-4-step-distilled-preview", }, + { + name: "VDN-H3", + modelIds: ["OpenVDN/vdn-minimax-h3"], + cookbook: "/cookbook/diffusion/MiniMax/MiniMax-H3#7-vdn-h3-hybrid-attention-8-step-distill", + }, { name: "MOVA", modelIds: ["OpenMOSS-Team/MOVA-360p", "OpenMOSS-Team/MOVA-720p"], diff --git a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh index 43a2b51b3..13ffbca31 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh @@ -47,8 +47,22 @@ struct QKNormRopePackKVParams : QKNormRopeParams { uint32_t suffix_tokens; }; -template -using QKNormRopeParamsT = std::conditional_t; +/// \brief Out-of-place variant: q/k are read (any strides), the normed + roped +/// rows are written to q_out/k_out and the inputs stay untouched. Used where a +/// consumer still needs the raw projections (VDN-H3's NoPE linear branch). +struct QKNormRopeOutOfPlaceParams : QKNormRopeParams { + void* __restrict__ q_out_ptr; + void* __restrict__ k_out_ptr; // pre-offset by -num_qo_heads * out_head_stride_bytes + int64_t q_out_stride_bytes; + int64_t k_out_stride_bytes; + int64_t out_head_stride_bytes; +}; + +template +using QKNormRopeParamsT = std::conditional_t< + kPackKV, + QKNormRopePackKVParams, + std::conditional_t>; constexpr uint32_t kThreadsPerBlock = 256; constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / device::kWarpThreads; @@ -199,9 +213,11 @@ template < bool kRoundNormBeforeRope, bool kPackKV, bool kCacheHasFullWidth, - typename IdType> -__global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT __grid_constant__ params) { + typename IdType, + bool kOutOfPlace = false> +__global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT __grid_constant__ params) { using namespace device; + static_assert(!(kPackKV && kOutOfPlace), "KV packing and out-of-place output are exclusive"); static_assert(std::is_same_v || std::is_same_v); static_assert(kHeadDim <= 256, "Only warp-level fused qknorm+rope is supported"); @@ -291,6 +307,13 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT __grid_c const void* input = load_q ? pointer::offset(q_ptr, token_id * q_stride_bytes, head_id * head_stride_bytes) : pointer::offset(k_ptr, token_id * k_stride_bytes, head_id * head_stride_bytes); void* output = const_cast(input); + if constexpr (kOutOfPlace) { + output = + load_q ? pointer::offset( + params.q_out_ptr, token_id * params.q_out_stride_bytes, head_id * params.out_head_stride_bytes) + : pointer::offset( + params.k_out_ptr, token_id * params.k_out_stride_bytes, head_id * params.out_head_stride_bytes); + } if constexpr (kPackKV) { if (!load_q) { const uint32_t batch_id = token_id / params.suffix_tokens; @@ -446,6 +469,23 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT __grid_c PDLTriggerSecondary(); } +/// \brief Shared launch tail of the three host runners: pick the index-type +/// instantiation, size the persistent grid from the occupancy table, launch. +template +void launch_qknorm_rope(const Params& params, bool is_int32, uint32_t num_works, DLDevice device, bool use_pdl) { + using namespace host; + const auto selected_kernel = is_int32 ? kKernelI32 : kKernelI64; + const uint32_t kNumSM = runtime::get_sm_count(device.device_id); + static const uint32_t kOccupancyTable[2] = { + runtime::get_blocks_per_sm(kKernelI32, kThreadsPerBlock), + runtime::get_blocks_per_sm(kKernelI64, kThreadsPerBlock), + }; + const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; + const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock); + const auto num_blocks = std::min(max_blocks, needed_blocks); + LaunchKernel(num_blocks, kThreadsPerBlock, device).enable_pdl(use_pdl)(selected_kernel, params); +} + template < int64_t kHeadDim, int64_t kRopeDim, @@ -528,18 +568,106 @@ struct QKNormRopeKernel { .eps = eps, }; - const auto is_int32 = id_type.is_type(); - const auto selected_kernel = is_int32 ? kernel : kernel; - const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); - static const uint32_t kOccupancyTable[2] = { - runtime::get_blocks_per_sm(kernel, kThreadsPerBlock), - runtime::get_blocks_per_sm(kernel, kThreadsPerBlock), - }; - const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; - const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock); - const auto num_blocks = std::min(max_blocks, needed_blocks); - LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params); + launch_qknorm_rope, kernel>( + params, id_type.is_type(), num_works, device.unwrap(), kUsePDL); + } +}; + +template < + int64_t kHeadDim, + int64_t kRopeDim, + bool kIsNeox, + bool kUsePDL, + typename DType, + typename CacheDType, + bool kRoundNormBeforeRope, + bool kCacheHasFullWidth> +struct QKNormRopeOutOfPlaceKernel { + static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported"); + template + static constexpr auto kernel = fused_qknorm_rope_warp< + kHeadDim, + kRopeDim, + kIsNeox, + kUsePDL, + DType, + CacheDType, + kRoundNormBeforeRope, + false, + kCacheHasFullWidth, + IdType, + true>; + + /// \brief QK-norm + RoPE from q/k into q_out/k_out; q and k are left untouched. + static void + run(const tvm::ffi::TensorView q, + const tvm::ffi::TensorView k, + const tvm::ffi::TensorView q_out, + const tvm::ffi::TensorView k_out, + const tvm::ffi::TensorView q_weight, + const tvm::ffi::TensorView k_weight, + const tvm::ffi::TensorView cos_sin_cache, + const tvm::ffi::TensorView positions, + float eps) { + using namespace host; + + auto N = SymbolicSize{"num_tokens"}; + auto Q = SymbolicSize{"num_qo_heads"}; + auto K = SymbolicSize{"num_kv_heads"}; + auto D = SymbolicSize{"head_dim"}; + auto Dq = SymbolicSize{"q_stride"}; + auto Dk = SymbolicSize{"k_stride"}; + auto Dd = SymbolicSize{"head_stride"}; + auto Dqo = SymbolicSize{"q_out_stride"}; + auto Dko = SymbolicSize{"k_out_stride"}; + auto Ddo = SymbolicSize{"out_head_stride"}; + auto device = SymbolicDevice{}; + auto id_type = SymbolicDType{}; + D.set_value(kHeadDim); + device.set_options(); + + TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype().with_device(device).verify(q); + TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype().with_device(device).verify(k); + TensorMatcher({N, Q, D}).with_strides({Dqo, Ddo, 1}).with_dtype().with_device(device).verify(q_out); + TensorMatcher({N, K, D}).with_strides({Dko, Ddo, 1}).with_dtype().with_device(device).verify(k_out); + TensorMatcher({D}).with_dtype().with_device(device).verify(q_weight).verify(k_weight); + TensorMatcher({-1, kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim}) + .with_dtype() + .with_device(device) + .verify(cos_sin_cache); + TensorMatcher({N}).with_dtype(id_type).with_device(device).verify(positions); + + const auto num_tokens = static_cast(N.unwrap()); + const auto num_qo_heads = static_cast(Q.unwrap()); + const auto num_kv_heads = static_cast(K.unwrap()); + if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return; + const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); + const auto out_head_stride_bytes = static_cast(Ddo.unwrap() * sizeof(DType)); + + QKNormRopeOutOfPlaceParams params{}; + params.q_ptr = q.data_ptr(); + params.k_ptr = pointer::offset(k.data_ptr(), -static_cast(num_qo_heads) * head_stride_bytes); + params.q_weight_ptr = q_weight.data_ptr(); + params.k_weight_ptr = k_weight.data_ptr(); + params.cos_sin_cache_ptr = cos_sin_cache.data_ptr(); + params.positions = positions.data_ptr(); + params.q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); + params.k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); + params.head_stride_bytes = head_stride_bytes; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.num_tokens = num_tokens; + params.eps = eps; + params.q_out_ptr = q_out.data_ptr(); + params.k_out_ptr = pointer::offset(k_out.data_ptr(), -static_cast(num_qo_heads) * out_head_stride_bytes); + params.q_out_stride_bytes = static_cast(Dqo.unwrap() * sizeof(DType)); + params.k_out_stride_bytes = static_cast(Dko.unwrap() * sizeof(DType)); + params.out_head_stride_bytes = out_head_stride_bytes; + + const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; + launch_qknorm_rope, kernel>( + params, id_type.is_type(), num_works, device.unwrap(), kUsePDL); } }; @@ -655,20 +783,11 @@ struct QKNormRopePackKVKernel { params.prefix_tokens = static_cast(prefix_tokens); params.suffix_tokens = static_cast(suffix_tokens); - const auto is_int32 = id_type.is_type(); - const auto selected_kernel = is_int32 ? kernel : kernel; - const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); - static const uint32_t kOccupancyTable[2] = { - runtime::get_blocks_per_sm(kernel, kThreadsPerBlock), - runtime::get_blocks_per_sm(kernel, kThreadsPerBlock), - }; - const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; const uint32_t num_prefix_works = static_cast(batch_size * prefix_tokens) * num_kv_heads; const uint32_t num_works = (num_qo_heads + num_kv_heads) * num_tokens + 2 * num_prefix_works + num_tokens * num_kv_heads; - const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock); - const auto num_blocks = std::min(max_blocks, needed_blocks); - LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params); + launch_qknorm_rope, kernel>( + params, id_type.is_type(), num_works, device.unwrap(), kUsePDL); } }; diff --git a/python/sglang/kernels/jit/csrc/diffusion/vdn_delta_factors.cuh b/python/sglang/kernels/jit/csrc/diffusion/vdn_delta_factors.cuh new file mode 100644 index 000000000..1a91cb4bd --- /dev/null +++ b/python/sglang/kernels/jit/csrc/diffusion/vdn_delta_factors.cuh @@ -0,0 +1,367 @@ +// SPDX-License-Identifier: Apache-2.0 +// Fused VDN-H3 delta-rule factors. +// +// Per (frame, head) the linear branch needs, for M = I + A (128x128 fp32, symmetric positive definite): +// transition = diag(alpha) M^-1 [F, H, dk, dk] +// injection = B M^-1 [F, H, dv, dk] +// The eager path is cholesky + solve_triangular + two GEMMs (~40 launches); this kernel is one +// launch: one CTA of 256 threads per matrix, thread (ti, tj) owns rows 8ti.., cols 8tj.. as float2 +// pairs t[8][4] so the rank-2 updates and the final GEMM run on packed FFMA2 (sm_100+, scalar +// fallback elsewhere). +// +// M^-1 is formed in place by block Gauss-Jordan elimination without pivoting (stable for SPD, the +// same class as Cholesky), two pivots per barrier. For the pivot block S = {k, k+1} with +// P = M[S,S], R = M[S,:], C = M[:,S], D the rest: +// R' = P^-1 R, G = C (raw), D' = D - G R', M'[~S, S] = 0 - G R'[:,S] = -G P^-1 +// (in-place GJ inverse semantics: the eliminated column receives the inverse column). The next +// block's band is prepared in the middle of the current update (software pipelined) and the barrier +// sits before the second half of the update so its FMAs overlap the barrier skew and the loads of +// the next step. Shared rows are column-swizzled so the float4 reads are bank-conflict free; the +// same swizzle is used for the smem copy of M^-1 consumed by the register-tiled GEMM B M^-1. +// +// Accuracy: the relative error vs fp64 matches the cholesky path; both are dominated by cond(M). +#include +#include + +#include + +#include + +namespace sglang { + +namespace vdn_delta_factors { + +namespace { + +constexpr int kDim = 128; // dk == dv == 128 +constexpr int kBlockSize = 256; // 16 x 16 tiles of 8 x 8 +constexpr int kMinBlocksPerSm = 2; // 128 registers per thread +constexpr int kXsBytes = kDim * kDim * static_cast(sizeof(float)); // 64 KB dynamic smem +constexpr unsigned kFullMask = 0xffffffffu; + +// thread tj's tile columns 8tj..8tj+3 land at 4tj.., 8tj+4..8tj+7 at 64+4tj.., so a quarter warp +// reads 8 consecutive 16-byte chunks +SGL_DEVICE int swz_lo(int tj) { + return 4 * tj; +} +SGL_DEVICE int swz_hi(int tj) { + return 64 + 4 * tj; +} + +SGL_DEVICE float rcp_nr(float p) { + float r; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(p)); + return fmaf(r, fmaf(-p, r, 1.f), r); // one Newton step: ~0.5 ulp +} +SGL_DEVICE float4 ld4(const float* p) { + return *reinterpret_cast(p); +} +SGL_DEVICE void st4(float* p, float a, float b, float c, float d) { + *reinterpret_cast(p) = make_float4(a, b, c, d); +} +SGL_DEVICE void st4(float* p, float2 a, float2 b) { + *reinterpret_cast(p) = make_float4(a.x, a.y, b.x, b.y); +} +SGL_DEVICE float2 f2(float a) { + return make_float2(a, a); +} +// packed fma: (a.x*b.x+c.x, a.y*b.y+c.y); one FFMA2 on sm_100+ (the first operand is a scalar +// broadcast in SASS), two FFMA elsewhere. Bitwise identical results either way. +SGL_DEVICE float2 fma2(float2 a, float2 b, float2 c) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + return __ffma2_rn(a, b, c); +#else + return make_float2(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y)); +#endif +} +SGL_DEVICE float get(const float4& v, int i) { + return i == 0 ? v.x : i == 1 ? v.y : i == 2 ? v.z : v.w; +} + +struct Smem { + float row[2][2][kDim]; // [buffer][band row m][swizzled col] scaled band rows R' + float col[2][kDim * 2]; // [buffer][row*2 + m] multipliers G +}; + +/// Prepare the band of pivot block {k, k+1}, k = 8*kt1 + R1 (R1 even). Column owners (tj == kt1) +/// publish G = C (0 for the band rows) and zero their band columns; row owners (ti == kt1) replace P +/// by I, scale R' = P^-1 R and publish it. Executed by all threads (P is broadcast with shuffles). +template +SGL_DEVICE void prepare(float2 (&t)[8][4], Smem& sm, int kt1, int ti, int tj, int i0) { + constexpr int CP = R1 >> 1; + const int src = ((kt1 & 1) << 4) | kt1; // lane of the diagonal tile inside the row-owner warp + const float pa = __shfl_sync(kFullMask, t[R1][CP].x, src); + const float pb = __shfl_sync(kFullMask, t[R1][CP].y, src); + const float pc = __shfl_sync(kFullMask, t[R1 + 1][CP].x, src); + const float pd = __shfl_sync(kFullMask, t[R1 + 1][CP].y, src); + if (tj == kt1) { + const bool diag = (ti == kt1); + float* dst = &sm.col[BUF][i0 * 2]; + float2 g[8]; +#pragma unroll + for (int r = 0; r < 8; ++r) { + const bool band = diag && (r == R1 || r == R1 + 1); + g[r] = band ? make_float2(0.f, 0.f) : t[r][CP]; + if (!band) t[r][CP] = make_float2(0.f, 0.f); + } +#pragma unroll + for (int q = 0; q < 4; ++q) + st4(dst + 4 * q, g[2 * q], g[2 * q + 1]); + } + if (ti == kt1) { + const float det = fmaf(pa, pd, -pb * pc); + const float idet = rcp_nr(det); + const float ia = pd * idet, ib = -pb * idet, ic = -pc * idet, id = pa * idet; + if (tj == kt1) { + t[R1][CP] = make_float2(1.f, 0.f); + t[R1 + 1][CP] = make_float2(0.f, 1.f); + } +#pragma unroll + for (int cp = 0; cp < 4; ++cp) { + const float2 x = t[R1][cp], y = t[R1 + 1][cp]; + t[R1][cp] = make_float2(fmaf(ia, x.x, ib * y.x), fmaf(ia, x.y, ib * y.y)); + t[R1 + 1][cp] = make_float2(fmaf(ic, x.x, id * y.x), fmaf(ic, x.y, id * y.y)); + } + st4(&sm.row[BUF][0][swz_lo(tj)], t[R1][0], t[R1][1]); + st4(&sm.row[BUF][0][swz_hi(tj)], t[R1][2], t[R1][3]); + st4(&sm.row[BUF][1][swz_lo(tj)], t[R1 + 1][0], t[R1 + 1][1]); + st4(&sm.row[BUF][1][swz_hi(tj)], t[R1 + 1][2], t[R1 + 1][3]); + } +} + +/// Rank-2 update for pivots (8kt + KR, 8kt + KR + 1); the next block's band is prepared in the middle. +template +SGL_DEVICE void step(float2 (&t)[8][4], Smem& sm, int kt, int ti, int tj, int i0) { + constexpr int CUR = (KR >> 1) & 1, NXT = CUR ^ 1; + constexpr int R1 = (KR + 2) & 7, CP = R1 >> 1; // next block's band rows / column pair + float2 rk0[4], rk1[4]; + float ck0[8], ck1[8]; // negated multipliers (scalar broadcast operands of FFMA2) + { + const float4 a = ld4(&sm.row[CUR][0][swz_lo(tj)]), b = ld4(&sm.row[CUR][0][swz_hi(tj)]); + const float4 c = ld4(&sm.row[CUR][1][swz_lo(tj)]), d = ld4(&sm.row[CUR][1][swz_hi(tj)]); + rk0[0] = make_float2(a.x, a.y); + rk0[1] = make_float2(a.z, a.w); + rk0[2] = make_float2(b.x, b.y); + rk0[3] = make_float2(b.z, b.w); + rk1[0] = make_float2(c.x, c.y); + rk1[1] = make_float2(c.z, c.w); + rk1[2] = make_float2(d.x, d.y); + rk1[3] = make_float2(d.z, d.w); + const float* cp = &sm.col[CUR][i0 * 2]; +#pragma unroll + for (int q = 0; q < 4; ++q) { + const float4 v = ld4(cp + 4 * q); + ck0[2 * q] = -v.x; + ck1[2 * q] = -v.y; + ck0[2 * q + 1] = -v.z; + ck1[2 * q + 1] = -v.w; + } + } + // part 1: the next band (rows R1, R1+1 fully; column pair CP in the other rows) +#pragma unroll + for (int cp = 0; cp < 4; ++cp) { + t[R1][cp] = fma2(f2(ck1[R1]), rk1[cp], fma2(f2(ck0[R1]), rk0[cp], t[R1][cp])); + t[R1 + 1][cp] = fma2(f2(ck1[R1 + 1]), rk1[cp], fma2(f2(ck0[R1 + 1]), rk0[cp], t[R1 + 1][cp])); + } +#pragma unroll + for (int r = 0; r < 8; ++r) { + if (r == R1 || r == R1 + 1) continue; + t[r][CP] = fma2(f2(ck1[r]), rk1[CP], fma2(f2(ck0[r]), rk0[CP], t[r][CP])); + } + if (KR != 6 || kt != 15) prepare(t, sm, kt + (KR == 6 ? 1 : 0), ti, tj, i0); + __syncthreads(); + // part 2: everything else (registers only; overlaps the barrier skew and the next step's loads) +#pragma unroll + for (int r = 0; r < 8; ++r) { + if (r == R1 || r == R1 + 1) continue; +#pragma unroll + for (int cp = 0; cp < 4; ++cp) { + if (cp == CP) continue; + t[r][cp] = fma2(f2(ck1[r]), rk1[cp], fma2(f2(ck0[r]), rk0[cp], t[r][cp])); + } + } +} + +/** + * \brief transition = diag(alpha) (I + A)^-1, injection = B (I + A)^-1 for a batch of 128x128 SPD A. + * + * \param A [N, 128, 128] fp32, symmetric positive semi-definite (I + A is inverted) + * \param B [N, 128, 128] fp32 + * \param alpha [N, 128] fp32 row scales of the transition + * \param transition [N, 128, 128] fp32 output + * \param injection [N, 128, 128] fp32 output + */ +__global__ void __launch_bounds__(kBlockSize, kMinBlocksPerSm) vdn_delta_factors_kernel( + const float* __restrict__ A, + const float* __restrict__ B, + const float* __restrict__ alpha, + float* __restrict__ transition, + float* __restrict__ injection) { + extern __shared__ __align__(16) float Xs[]; // [kDim][kDim], column-swizzled copy of (I + A)^-1 + __shared__ __align__(16) Smem sm; + + const int n = blockIdx.x; + const int tid = threadIdx.x; + const int ti = tid >> 4, tj = tid & 15; + const int i0 = ti * 8, j0 = tj * 8; + const size_t mat = static_cast(n) * kDim * kDim; + const float* An = A + mat; + const float* Bn = B + mat; + + float2 t[8][4]; +#pragma unroll + for (int r = 0; r < 8; ++r) { + const float4 v0 = ld4(An + (i0 + r) * kDim + j0), v1 = ld4(An + (i0 + r) * kDim + j0 + 4); + t[r][0] = make_float2(v0.x, v0.y); + t[r][1] = make_float2(v0.z, v0.w); + t[r][2] = make_float2(v1.x, v1.y); + t[r][3] = make_float2(v1.z, v1.w); + } + if (ti == tj) { // M = I + A +#pragma unroll + for (int r = 0; r < 8; ++r) { + if (r & 1) { + t[r][r >> 1].y += 1.f; + } else { + t[r][r >> 1].x += 1.f; + } + } + } + + prepare<0, 0>(t, sm, 0, ti, tj, i0); + __syncthreads(); +#pragma unroll 1 + for (int kt = 0; kt < kDim / 8; ++kt) { + step<0>(t, sm, kt, ti, tj, i0); + step<2>(t, sm, kt, ti, tj, i0); + step<4>(t, sm, kt, ti, tj, i0); + step<6>(t, sm, kt, ti, tj, i0); + } + + // transition = diag(alpha) X, straight from registers + { + const float* al = alpha + static_cast(n) * kDim + i0; + const float4 a0 = ld4(al), a1 = ld4(al + 4); + const float av[8] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w}; + float* Tn = transition + mat; +#pragma unroll + for (int r = 0; r < 8; ++r) { + float* row = Tn + (i0 + r) * kDim + j0; + st4(row, av[r] * t[r][0].x, av[r] * t[r][0].y, av[r] * t[r][1].x, av[r] * t[r][1].y); + st4(row + 4, av[r] * t[r][2].x, av[r] * t[r][2].y, av[r] * t[r][3].x, av[r] * t[r][3].y); + } + } + // stage X in smem (swizzled columns) for the GEMM +#pragma unroll + for (int r = 0; r < 8; ++r) { + float* row = Xs + (i0 + r) * kDim; + st4(row + swz_lo(tj), t[r][0], t[r][1]); + st4(row + swz_hi(tj), t[r][2], t[r][3]); + } + __syncthreads(); + + // injection = B X (register-tiled GEMM; B rows from global through L1, X from smem) +#pragma unroll + for (int r = 0; r < 8; ++r) +#pragma unroll + for (int cp = 0; cp < 4; ++cp) + t[r][cp] = make_float2(0.f, 0.f); + const float* Bp = Bn + i0 * kDim; +#pragma unroll 1 + for (int k = 0; k < kDim; k += 4) { + float4 b[8]; +#pragma unroll + for (int r = 0; r < 8; ++r) + b[r] = __ldg(reinterpret_cast(Bp + r * kDim + k)); +#pragma unroll + for (int kk = 0; kk < 4; ++kk) { + const float* xr = Xs + (k + kk) * kDim; + const float4 x0 = ld4(xr + swz_lo(tj)), x1 = ld4(xr + swz_hi(tj)); + const float2 xv[4] = { + make_float2(x0.x, x0.y), make_float2(x0.z, x0.w), make_float2(x1.x, x1.y), make_float2(x1.z, x1.w)}; +#pragma unroll + for (int r = 0; r < 8; ++r) { + const float2 bv = f2(get(b[r], kk)); +#pragma unroll + for (int cp = 0; cp < 4; ++cp) + t[r][cp] = fma2(bv, xv[cp], t[r][cp]); + } + } + } + { + float* Jn = injection + mat; +#pragma unroll + for (int r = 0; r < 8; ++r) { + float* row = Jn + (i0 + r) * kDim + j0; + st4(row, t[r][0], t[r][1]); + st4(row + 4, t[r][2], t[r][3]); + } + } +} + +} // namespace + +struct VdnDeltaFactorsKernel { + /** + * \brief Validate the tensors and launch one CTA per matrix. + * + * \param transition [N, 128, 128] fp32 output, diag(alpha) (I + A)^-1 + * \param injection [N, 128, 128] fp32 output, B (I + A)^-1 + * \param A [N, 128, 128] fp32 SPD statistics (I + A is inverted) + * \param B [N, 128, 128] fp32 + * \param alpha [N, 128] fp32 + */ + static void + run(tvm::ffi::TensorView transition, + tvm::ffi::TensorView injection, + tvm::ffi::TensorView A, + tvm::ffi::TensorView B, + tvm::ffi::TensorView alpha) { + using namespace host; + auto N = SymbolicSize{"num_matrices"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({N, kDim, kDim}) + .with_dtype() + .with_device(device) + .verify(transition) + .verify(injection) + .verify(A) + .verify(B); + TensorMatcher({N, kDim}).with_dtype().with_device(device).verify(alpha); + const int64_t num = N.unwrap(); + if (num == 0) return; + CHECK_HOST( + transition.data_ptr() != A.data_ptr() && transition.data_ptr() != B.data_ptr() && + injection.data_ptr() != A.data_ptr() && injection.data_ptr() != B.data_ptr() && + transition.data_ptr() != injection.data_ptr()) + << "vdn_delta_factors outputs must not alias inputs"; + // every tensor is read or written as float4; a storage offset breaks this + const auto aligned16 = [](const void* p) { return reinterpret_cast(p) % 16 == 0; }; + CHECK_HOST( + aligned16(A.data_ptr()) && aligned16(B.data_ptr()) && aligned16(alpha.data_ptr()) && + aligned16(transition.data_ptr()) && aligned16(injection.data_ptr())) + << "vdn_delta_factors needs 16-byte aligned tensors"; + const DLDevice dev = device.unwrap(); + // 64 KB of dynamic shared memory needs the opt-in, once per device. + static bool attr_set[64] = {}; + const int dev_id = dev.device_id; + CHECK_HOST(dev_id >= 0 && dev_id < 64) << "vdn_delta_factors: unexpected device id " << dev_id; + if (!attr_set[dev_id]) { + CHECK_CUDA(cudaFuncSetAttribute(vdn_delta_factors_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kXsBytes)) + << "vdn_delta_factors: cannot reserve " << kXsBytes << " bytes of dynamic shared memory"; + attr_set[dev_id] = true; + } + LaunchKernel(static_cast(num), kBlockSize, dev, kXsBytes)( + vdn_delta_factors_kernel, + static_cast(A.data_ptr()), + static_cast(B.data_ptr()), + static_cast(alpha.data_ptr()), + static_cast(transition.data_ptr()), + static_cast(injection.data_ptr())); + } +}; + +} // namespace vdn_delta_factors + +} // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 8bc44770b..4aeb0c87e 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/python/sglang/kernels/ops/diffusion/README.md @@ -37,6 +37,7 @@ norm/ RMSNorm / LayerNorm / GroupNorm and their fused epilogues modulate/ adaLN modulate, gating, timestep conditioning rope/ rotary embeddings and the QK-norm chains fused into them activation/ SiLU / GLU / GELU fusions +quantization/ MXFP8 producers whose scales land in the GEMM's swizzled layout attention/ sparse linear attention, gated delta-net routing/ diffusion-model MoE routing and expert selection layout/ pure data movement: USP/Ulysses relayout, varlen pack, causal pad @@ -140,6 +141,7 @@ tensor copy per residual site. |---|---|---| | `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs split baseline; `round_norm_before_rope=True` makes it exact; supports compact and full-width NeoX/interleaved caches | | `fused_qknorm_rope_pack_kv` | JIT CUDA | as above, also packs prefix K/V | +| `fused_qknorm_rope_out_of_place` | JIT CUDA | as above, bit-equal to the in-place kernel; reads strided q/k and writes contiguous copies, inputs untouched (VDN-H3 keeps the raw q/k for its linear branch) | | `try_fused_flux2_qkv_epilogue` | KDA (JIT CUDA) | bit-exact vs the selected BF16 chain | FLUX.2 QK RMSNorm + RoPE + joint QKV packing | | `try_fused_qwen_qkv_epilogue` | JIT CUDA | bit-exact vs the selected BF16 chain | Qwen-Image QK RMSNorm + RoPE + joint QKV writes; SM100+ | | `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) | @@ -150,6 +152,22 @@ tensor copy per residual site. | `apply_rotary_embedding` | Triton (+fallbacks) | close; the generic entry point | | `hunyuan_qkv_rope_pack` | Triton | bit-exact; packs QKV and applies RoPE in one pass | +### MiniMax-H3 / VDN-H3 linear branch + +| Entry point | Backend | Contract | +|---|---|---| +| `vdn_frame_stats_prep`, `vdn_gather_linear_state` | Triton | bit-exact (same products, fp32 gather) | +| `vdn_temporal_conv_act`, `vdn_silu_l2norm`, `vdn_linear_epilogue` | Triton | one rounding at the store, within one bf16 ulp of the eager chain; the model's own inference kernels, mounted unconditionally by the VDN-H3 branch | +| `vdn_delta_factors` | JIT CUDA | `(alpha * inv(I + A), B @ inv(I + A))` in one launch; same fp32 accuracy class as the cholesky + solve_triangular chain (cond-dominated); head_dim 128 | + +### MXFP8 producers (online `mxfp8`, cuBLASLt block-scaled GEMM on SM100) + +| Entry point | Backend | Contract | +|---|---|---| +| `mxfp8_quantize_swizzled` | Triton | bit-exact vs `flashinfer.mxfp8_quantize(x, True)`: e4m3 payload + block-32 E8M0 scales in the `SWIZZLE_32_4_4` layout; weights at load and any bf16 GEMM input | +| `silu_mul_mxfp8` | Triton | bit-exact vs eager bf16 `silu(gate) * up` followed by the quantizer above; the fc2 input | +| `indexed_scale_shift_mxfp8_` | Triton | bit-exact vs `indexed_scale_shift_bf16_` followed by the quantizer above, optionally keeping the bf16 rows in place; the qkv / fc1 inputs | + ### MoE routing | Entry point | Backend | Contract | Applies to | diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 38668208e..f584fd422 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -328,6 +328,76 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _CUDA, "Sana-WM bidirectional gated delta-net.", ), + ( + "diffusion.fused_qknorm_rope_out_of_place", + KernelBackend.JIT, + "rope.qknorm_rope_jit:fused_qknorm_rope_out_of_place", + _CUDA, + "Out-of-place fused QK-norm + RoPE (raw q/k preserved).", + ), + ( + "diffusion.vdn_delta_factors", + KernelBackend.JIT, + "attention.vdn_delta_factors_jit:vdn_delta_factors", + _CUDA, + "VDN-H3 delta rule: fused (I + A)^-1 -> transition / injection (fp32, head_dim 128).", + ), + ( + "diffusion.vdn_temporal_conv_act", + KernelBackend.TRITON, + "attention.vdn_linear_branch_triton:vdn_temporal_conv_act", + _CUDA, + "VDN-H3 linear branch: 5-tap temporal conv + SiLU + L2 norm.", + ), + ( + "diffusion.vdn_silu_l2norm", + KernelBackend.TRITON, + "attention.vdn_linear_branch_triton:vdn_silu_l2norm", + _CUDA, + "VDN-H3 linear branch: SiLU + L2 norm over head_dim.", + ), + ( + "diffusion.vdn_frame_stats_prep", + KernelBackend.TRITON, + "attention.vdn_linear_branch_triton:vdn_frame_stats_prep", + _CUDA, + "VDN-H3 linear branch: frame-statistics GEMM operands in one pass.", + ), + ( + "diffusion.vdn_gather_linear_state", + KernelBackend.TRITON, + "attention.vdn_linear_branch_triton:vdn_gather_linear_state", + _CUDA, + "VDN-H3 linear branch: alpha-bridged boundary gather in one pass.", + ), + ( + "diffusion.vdn_linear_epilogue", + KernelBackend.TRITON, + "attention.vdn_linear_branch_triton:vdn_linear_epilogue", + _CUDA, + "VDN-H3 linear branch: RMSNorm * gate readout epilogue.", + ), + ( + "diffusion.mxfp8_quantize_swizzled", + KernelBackend.TRITON, + "quantization.mxfp8_swizzled_triton:mxfp8_quantize_swizzled", + _CUDA, + "bf16 -> MXFP8 (e4m3, block-32 E8M0 scales in the cuBLASLt swizzled layout).", + ), + ( + "diffusion.silu_mul_mxfp8", + KernelBackend.TRITON, + "quantization.mxfp8_swizzled_triton:silu_mul_mxfp8", + _CUDA, + "SwiGLU + MXFP8 quant for the online mxfp8 fc2 input.", + ), + ( + "diffusion.indexed_scale_shift_mxfp8_", + KernelBackend.TRITON, + "quantization.mxfp8_swizzled_triton:indexed_scale_shift_mxfp8_", + _CUDA, + "Indexed adaLN modulation + MXFP8 quant for the online mxfp8 qkv/fc1 inputs.", + ), ( "diffusion.group_limited_topk", KernelBackend.TRITON, @@ -541,6 +611,24 @@ _EXPORTS: dict[str, str] = { "fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit", "fused_causal_conv3d_cat_pad": "layout.causal_conv3d_cat_pad_triton", "pack_qkv_destination_major": "layout.ulysses_qkv_triton", + "fused_qknorm_rope_out_of_place": "rope.qknorm_rope_jit", + "vdn_delta_factors": "attention.vdn_delta_factors_jit", + "can_use_vdn_delta_factors": "attention.vdn_delta_factors_jit", + "vdn_temporal_conv_act": "attention.vdn_linear_branch_triton", + "can_use_vdn_temporal_conv_act": "attention.vdn_linear_branch_triton", + "can_use_vdn_silu_l2norm": "attention.vdn_linear_branch_triton", + "can_use_vdn_frame_stats_prep": "attention.vdn_linear_branch_triton", + "can_use_vdn_gather_linear_state": "attention.vdn_linear_branch_triton", + "can_use_vdn_linear_epilogue": "attention.vdn_linear_branch_triton", + "vdn_silu_l2norm": "attention.vdn_linear_branch_triton", + "vdn_frame_stats_prep": "attention.vdn_linear_branch_triton", + "vdn_gather_linear_state": "attention.vdn_linear_branch_triton", + "vdn_linear_epilogue": "attention.vdn_linear_branch_triton", + "can_use_mxfp8_swizzled": "quantization.mxfp8_swizzled_triton", + "can_use_silu_mul_mxfp8": "quantization.mxfp8_swizzled_triton", + "indexed_scale_shift_mxfp8_": "quantization.mxfp8_swizzled_triton", + "mxfp8_quantize_swizzled": "quantization.mxfp8_swizzled_triton", + "silu_mul_mxfp8": "quantization.mxfp8_swizzled_triton", "can_use_usp_merge_heads": "layout.usp_relayout_jit", "usp_merge_heads": "layout.usp_relayout_jit", "build_inv_indices": "layout.varlen_pack_pad_triton", diff --git a/python/sglang/kernels/ops/diffusion/attention/vdn_delta_factors_jit.py b/python/sglang/kernels/ops/diffusion/attention/vdn_delta_factors_jit.py new file mode 100644 index 000000000..4afd2b4ef --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/attention/vdn_delta_factors_jit.py @@ -0,0 +1,102 @@ +"""Fused VDN-H3 delta-rule factors: (I + A)^-1 folded into the transition and injection. + +One CUDA kernel (block Gauss-Jordan inverse in registers + the two products) replaces the +cholesky / solve_triangular / GEMM chain of ``delta_factor_apply`` for the ``vdn_solve`` and +``vdn_scaled`` rules. fp32, head_dim 128 only. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +HEAD_DIM = 128 +_FLOAT4_BYTES = 16 + + +@cache_once +def _jit_vdn_delta_factors_module() -> Module: + if torch.cuda.get_device_capability()[0] < 8: + raise RuntimeError( + "vdn_delta_factors needs SM80 or later (2 x 70 KB shared memory per SM)" + ) + return load_jit( + "diffusion_vdn_delta_factors", + cuda_files=["diffusion/vdn_delta_factors.cuh"], + cuda_wrappers=[ + ("vdn_delta_factors", "vdn_delta_factors::VdnDeltaFactorsKernel::run") + ], + ) + + +def _aligned(t: torch.Tensor) -> torch.Tensor: + # the kernel loads float4; .contiguous() keeps a storage offset, a fresh allocation is aligned + return t if t.data_ptr() % _FLOAT4_BYTES == 0 else t.clone() + + +def _fake_impl( + A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + del alpha + return torch.empty_like(A), torch.empty_like(B) + + +@register_custom_op( + op_name="diffusion_vdn_delta_factors", + mutates_args=[], + fake_impl=_fake_impl, +) +def vdn_delta_factors( + A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """``(alpha[..., :, None] * inv(I + A), B @ inv(I + A))`` for fp32 ``[..., 128, 128]`` SPD ``A``. + + ``B`` has the shape of ``A``; ``alpha`` is ``[..., 128]``. Same accuracy as the eager + cholesky path (both are dominated by cond(I + A) in fp32). + """ + A, B, alpha = _aligned(A), _aligned(B), _aligned(alpha) + transition = torch.empty_like(A) + injection = torch.empty_like(B) + module = _jit_vdn_delta_factors_module() + module.vdn_delta_factors( + transition.view(-1, HEAD_DIM, HEAD_DIM), + injection.view(-1, HEAD_DIM, HEAD_DIM), + A.view(-1, HEAD_DIM, HEAD_DIM), + B.view(-1, HEAD_DIM, HEAD_DIM), + alpha.view(-1, HEAD_DIM), + ) + return transition, injection + + +def can_use_vdn_delta_factors( + A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor +) -> bool: + return ( + A.is_cuda + and A.dtype is torch.float32 + and B.dtype is torch.float32 + and alpha.dtype is torch.float32 + and A.device == B.device == alpha.device + and A.dim() >= 2 + and A.shape[-1] == HEAD_DIM + and A.shape[-2] == HEAD_DIM + and B.shape == A.shape + and alpha.shape == A.shape[:-1] + and A.is_contiguous() + and B.is_contiguous() + and alpha.is_contiguous() + and torch.cuda.get_device_capability(A.device)[0] >= 8 + ) + + +__all__ = [ + "can_use_vdn_delta_factors", + "vdn_delta_factors", +] diff --git a/python/sglang/kernels/ops/diffusion/attention/vdn_linear_branch_triton.py b/python/sglang/kernels/ops/diffusion/attention/vdn_linear_branch_triton.py new file mode 100644 index 000000000..d27c08990 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/attention/vdn_linear_branch_triton.py @@ -0,0 +1,569 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused Triton kernels for the VDN-H3 (Video DeltaNet MiniMax-H3) linear branch; +each reads its operands once and rounds once at the store. + + vdn_temporal_conv_act 5-tap depthwise temporal conv + SiLU [+ L2 norm] + (port of OpenVDN's _tconv_act_kernel) + vdn_silu_l2norm SiLU [+ L2 norm] over head_dim, strided input ok + vdn_frame_stats_prep the four GEMM operands of the frame statistics + (kf16, kf32, kf32 * beta, v * beta) in [F, H, S, d] + off one read of k and one of v + vdn_gather_linear_state the alpha-bridged boundary gather over the fp32 + state banks + vdn_linear_epilogue RMSNorm(d) * gate with the [F, H, S, d] -> [F*S, H*d] + transpose folded into the store + +Contract: vdn_frame_stats_prep and vdn_gather_linear_state are bitwise equal +to the eager chains (widening casts, same products, fp32 gather). The three +activation kernels round once instead of once per op and sit within one bf16 +ulp of the eager bf16 chains; OpenVDN ships the same contract, so the branch +mounts them unconditionally. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +_BLOCK_T = 16 +_BLOCK_ROWS = 32 + + +def _pow2_head_dim(head_dim: int) -> bool: + return head_dim >= 16 and head_dim & (head_dim - 1) == 0 + + +def _check_head_dim(head_dim: int) -> None: + if not _pow2_head_dim(head_dim): + raise ValueError( + f"head_dim must be a power of two >= 16 (tl.arange), got {head_dim}" + ) + + +def _cuda_bf16_rows(t: torch.Tensor) -> bool: + return t.is_cuda and t.dtype == torch.bfloat16 and t.stride(-1) == 1 + + +def _i32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.int32).contiguous() + + +def can_use_vdn_temporal_conv_act(x: torch.Tensor, heads: int, head_dim: int) -> bool: + """x [T, S, heads * head_dim] bf16 on CUDA, power-of-two head_dim.""" + return ( + _cuda_bf16_rows(x) + and x.ndim == 3 + and x.shape[-1] == heads * head_dim + and _pow2_head_dim(head_dim) + and not torch.compiler.is_compiling() + ) + + +def can_use_vdn_silu_l2norm(tokens: torch.Tensor) -> bool: + """tokens [N, H, d] bf16 on CUDA (any row/head strides), power-of-two d.""" + return ( + _cuda_bf16_rows(tokens) + and tokens.ndim == 3 + and _pow2_head_dim(tokens.shape[-1]) + and not torch.compiler.is_compiling() + ) + + +def can_use_vdn_frame_stats_prep(key: torch.Tensor, value: torch.Tensor) -> bool: + """key/value [F * S, H, d] bf16 on CUDA with matching shapes.""" + return ( + _cuda_bf16_rows(key) + and _cuda_bf16_rows(value) + and key.ndim == 3 + and key.shape == value.shape + and _pow2_head_dim(key.shape[-1]) + and not torch.compiler.is_compiling() + ) + + +def can_use_vdn_gather_linear_state(prefix: torch.Tensor) -> bool: + """prefix/suffix [F, H, dv, dk] fp32 on CUDA, power-of-two dk.""" + return ( + prefix.is_cuda + and prefix.dtype == torch.float32 + and prefix.ndim == 4 + and _pow2_head_dim(prefix.shape[-1]) + and not torch.compiler.is_compiling() + ) + + +def can_use_vdn_linear_epilogue(readout: torch.Tensor) -> bool: + """readout [F, H, S, d] bf16 on CUDA, power-of-two d.""" + return ( + _cuda_bf16_rows(readout) + and readout.ndim == 4 + and _pow2_head_dim(readout.shape[-1]) + and not torch.compiler.is_compiling() + ) + + +# -------------------------------------------------------------------------- +# temporal conv + SiLU + L2 norm +# -------------------------------------------------------------------------- + + +@triton.jit +def _tconv_act_kernel( + X, + W, + OUT, + num_frames, + tokens_per_frame, + channels, + BLOCK_T: tl.constexpr, + HEAD_DIM: tl.constexpr, + L2NORM: tl.constexpr, + HEADS: tl.constexpr, + FRAME_MAJOR: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_s = tl.program_id(1) + pid_h = tl.program_id(2) + chan = pid_h * HEAD_DIM + tl.arange(0, HEAD_DIM) + rows = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) + valid = rows < num_frames + + acc = tl.zeros((BLOCK_T, HEAD_DIM), dtype=tl.float32) + for dt in tl.static_range(5): + r = rows + dt - 2 + ok = valid & (r >= 0) & (r < num_frames) # zero padding, both ends + v = tl.load( + X + + (r[:, None].to(tl.int64) * tokens_per_frame + pid_s) * channels + + chan[None, :], + mask=ok[:, None], + other=0.0, + ).to(tl.float32) + wd = tl.load(W + chan * 5 + dt).to(tl.float32) + acc += v * wd[None, :] + + y = acc * tl.sigmoid(acc) # SiLU + if L2NORM: + inv = 1.0 / tl.sqrt(tl.maximum(tl.sum(y * y, axis=1), 1e-12)) + y = y * inv[:, None] + if FRAME_MAJOR: + # [T, HEADS, S, D]: the readout bmm reads this layout directly + dst = ( + (rows[:, None].to(tl.int64) * HEADS + pid_h) * tokens_per_frame + pid_s + ) * HEAD_DIM + tl.arange(0, HEAD_DIM)[None, :] + else: + dst = (rows[:, None].to(tl.int64) * tokens_per_frame + pid_s) * channels + chan[ + None, : + ] + tl.store(OUT + dst, y.to(OUT.dtype.element_ty), mask=valid[:, None]) + + +def vdn_temporal_conv_act( + x: torch.Tensor, + w: torch.Tensor, + heads: int, + head_dim: int, + l2norm: bool, + frame_major: bool = False, +) -> torch.Tensor: + """x [T, S, C] bf16 contiguous, w [C, 5] -> [T * S, heads, head_dim], or + [T, heads, S, head_dim] with ``frame_major``.""" + if not x.is_cuda: + raise ValueError("vdn_temporal_conv_act is a Triton kernel; x must be on CUDA") + _check_head_dim(head_dim) + num_frames, tokens_per_frame, channels = x.shape + if channels != heads * head_dim: + raise ValueError(f"C={channels} != heads*head_dim={heads * head_dim}") + if w.shape != (channels, 5): + raise ValueError(f"w must be [C, 5], got {tuple(w.shape)}") + x = x.contiguous() + w = w.contiguous() + out = torch.empty_like(x) + _tconv_act_kernel[(triton.cdiv(num_frames, _BLOCK_T), tokens_per_frame, heads)]( + x, + w, + out, + num_frames, + tokens_per_frame, + channels, + BLOCK_T=_BLOCK_T, + HEAD_DIM=head_dim, + L2NORM=l2norm, + HEADS=heads, + FRAME_MAJOR=frame_major, + num_warps=4, + num_stages=2, + ) + if frame_major: + return out.view(num_frames, heads, tokens_per_frame, head_dim) + return out.view(num_frames * tokens_per_frame, heads, head_dim) + + +# -------------------------------------------------------------------------- +# SiLU + L2 norm on a (possibly strided) [N, H, d] tensor +# -------------------------------------------------------------------------- + + +@triton.jit +def _silu_l2norm_kernel( + X, + OUT, + N, + stride_n, + stride_h, + H, + tokens_per_frame, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + L2NORM: tl.constexpr, + FRAME_MAJOR: tl.constexpr, +): + pid_n = tl.program_id(0) + pid_h = tl.program_id(1) + rows = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + valid = rows < N + offs = tl.arange(0, HEAD_DIM) + x = tl.load( + X + rows[:, None].to(tl.int64) * stride_n + pid_h * stride_h + offs[None, :], + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + y = x * tl.sigmoid(x) + if L2NORM: + inv = 1.0 / tl.sqrt(tl.maximum(tl.sum(y * y, axis=1), 1e-12)) + y = y * inv[:, None] + if FRAME_MAJOR: + # row n = frame * S_ + s -> [F, H, S_, D] + frame = rows // tokens_per_frame + pos = rows - frame * tokens_per_frame + dst = ( + (frame[:, None].to(tl.int64) * H + pid_h) * tokens_per_frame + pos[:, None] + ) * HEAD_DIM + offs[None, :] + else: + dst = (rows[:, None].to(tl.int64) * H + pid_h) * HEAD_DIM + offs[None, :] + tl.store(OUT + dst, y.to(OUT.dtype.element_ty), mask=valid[:, None]) + + +def vdn_silu_l2norm( + tokens: torch.Tensor, l2norm: bool, per_frame: int | None = None +) -> torch.Tensor: + """tokens [N, H, d] (last dim contiguous) -> contiguous [N, H, d], or + [N / per_frame, H, per_frame, d] when ``per_frame`` is given.""" + if not tokens.is_cuda: + raise ValueError("vdn_silu_l2norm is a Triton kernel; tokens must be on CUDA") + N, H, D = tokens.shape + _check_head_dim(D) + if tokens.stride(-1) != 1: + tokens = tokens.contiguous() + frame_major = per_frame is not None + if frame_major and (per_frame <= 0 or N % per_frame): + raise ValueError(f"per_frame={per_frame} must divide N={N}") + shape = (N // per_frame, H, per_frame, D) if frame_major else (N, H, D) + out = torch.empty(shape, dtype=tokens.dtype, device=tokens.device) + if N == 0: + return out + _silu_l2norm_kernel[(triton.cdiv(N, _BLOCK_ROWS), H)]( + tokens, + out, + N, + tokens.stride(0), + tokens.stride(1), + H, + per_frame if frame_major else 1, + BLOCK_N=_BLOCK_ROWS, + HEAD_DIM=D, + L2NORM=l2norm, + FRAME_MAJOR=frame_major, + num_warps=4, + ) + return out + + +# -------------------------------------------------------------------------- +# frame statistics prologue +# -------------------------------------------------------------------------- + + +@triton.jit +def _frame_stats_prep_kernel( + K, + V, + BETA, + K16, + K32, + KB32, + VB, + tokens_per_frame, + H, + BLOCK_S: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + pid_s = tl.program_id(0) + f = tl.program_id(1) + h = tl.program_id(2) + s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + valid = s < tokens_per_frame + offs = tl.arange(0, HEAD_DIM) + rows = f * tokens_per_frame + s # token rows + src = (rows[:, None].to(tl.int64) * H + h) * HEAD_DIM + offs[None, :] # [F*S, H, d] + dst = ((f * H + h) * tokens_per_frame + s)[:, None] * HEAD_DIM + offs[ + None, : + ] # [F, H, S, d] + k = tl.load(K + src, mask=valid[:, None], other=0.0) + v = tl.load(V + src, mask=valid[:, None], other=0.0) + beta = tl.load(BETA + rows * H + h, mask=valid, other=0.0) + k32 = k.to(tl.float32) + beta32 = beta.to(tl.float32) + tl.store(K16 + dst, k, mask=valid[:, None]) + tl.store(K32 + dst, k32, mask=valid[:, None]) + tl.store(KB32 + dst, k32 * beta32[:, None], mask=valid[:, None]) + vb = (v.to(tl.float32) * beta32[:, None]).to(VB.dtype.element_ty) + tl.store(VB + dst, vb, mask=valid[:, None]) + + +def vdn_frame_stats_prep( + key: torch.Tensor, + value: torch.Tensor, + beta: torch.Tensor, + num_frames: int, + tokens_per_frame: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """key/value [F*S, H, d] bf16 contiguous, beta [F*S, H] bf16 -> + (k16 [F,H,S,d] bf16, k32 fp32, k32*beta fp32, v*beta bf16), all contiguous.""" + if not key.is_cuda: + raise ValueError( + "vdn_frame_stats_prep is a Triton kernel; inputs must be on CUDA" + ) + rows, H, D = key.shape + _check_head_dim(D) + if rows != num_frames * tokens_per_frame: + raise ValueError(f"{rows} rows != {num_frames} x {tokens_per_frame}") + key = key.contiguous() + value = value.contiguous() + beta = beta.to(key.dtype).contiguous() + shape = (num_frames, H, tokens_per_frame, D) + k16 = torch.empty(shape, dtype=key.dtype, device=key.device) + k32 = torch.empty(shape, dtype=torch.float32, device=key.device) + kb32 = torch.empty(shape, dtype=torch.float32, device=key.device) + vb = torch.empty(shape, dtype=value.dtype, device=key.device) + _frame_stats_prep_kernel[ + (triton.cdiv(tokens_per_frame, _BLOCK_ROWS), num_frames, H) + ]( + key, + value, + beta, + k16, + k32, + kb32, + vb, + tokens_per_frame, + H, + BLOCK_S=_BLOCK_ROWS, + HEAD_DIM=D, + num_warps=4, + ) + return k16, k32, kb32, vb + + +# -------------------------------------------------------------------------- +# readout epilogue: RMSNorm(d) * gate, [F, H, S, d] -> [F*S, H*d] +# -------------------------------------------------------------------------- + + +@triton.jit +def _linear_epilogue_kernel( + R, + W, + G, + OUT, + tokens_per_frame, + H, + eps, + BLOCK_S: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + pid_s = tl.program_id(0) + f = tl.program_id(1) + h = tl.program_id(2) + s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + valid = s < tokens_per_frame + offs = tl.arange(0, HEAD_DIM) + src = ((f * H + h) * tokens_per_frame + s)[:, None] * HEAD_DIM + offs[None, :] + rows = f * tokens_per_frame + s + dst = (rows[:, None].to(tl.int64) * H + h) * HEAD_DIM + offs[None, :] + r = tl.load(R + src, mask=valid[:, None], other=0.0).to(tl.float32) + ms = tl.sum(r * r, axis=1) / HEAD_DIM + w = tl.load(W + offs).to(tl.float32) + g = tl.load(G + dst, mask=valid[:, None], other=0.0).to(tl.float32) + y = r * (1.0 / tl.sqrt(ms + eps))[:, None] * w[None, :] * g + tl.store(OUT + dst, y.to(OUT.dtype.element_ty), mask=valid[:, None]) + + +def vdn_linear_epilogue( + readout: torch.Tensor, + norm_weight: torch.Tensor, + gate: torch.Tensor, + eps: float, +) -> torch.Tensor: + """readout [F, H, S, d] bf16 contiguous, norm_weight [d], gate [F*S, H, d] + -> [F*S, H*d] bf16.""" + if not readout.is_cuda: + raise ValueError( + "vdn_linear_epilogue is a Triton kernel; readout must be on CUDA" + ) + F, H, tokens_per_frame, D = readout.shape + _check_head_dim(D) + readout = readout.contiguous() + gate = gate.reshape(F * tokens_per_frame, H, D).to(readout.dtype).contiguous() + out = torch.empty( + (F * tokens_per_frame, H * D), dtype=readout.dtype, device=readout.device + ) + _linear_epilogue_kernel[(triton.cdiv(tokens_per_frame, _BLOCK_ROWS), F, H)]( + readout, + norm_weight.contiguous(), + gate, + out, + tokens_per_frame, + H, + float(eps), + BLOCK_S=_BLOCK_ROWS, + HEAD_DIM=D, + num_warps=4, + ) + return out + + +# -------------------------------------------------------------------------- +# boundary gather: prefix[lo-1] * prod alpha + suffix[hi+1] * prod alpha +# -------------------------------------------------------------------------- + + +@triton.jit +def _gather_state_kernel( + PREFIX, + SUFFIX, + LOGP, # [F+1, H, dk] fp32 exclusive log-alpha prefix sums + TEXT, # [H, dv, dk] fp32 (or PREFIX when HAS_TEXT is False) + BEFORE, # [F] int32 prefix row to read (clamped) + AFTER, # [F] int32 suffix row to read (clamped) + HASB, # [F] int32 0/1 + HASA, # [F] int32 0/1 + BRIDGEB, # [F] int32 log-prefix row for the before side + BRIDGEA, # [F] int32 log-prefix row for the after side + OUT, + H, + DV, + HAS_TEXT: tl.constexpr, + BRIDGE: tl.constexpr, + BLOCK_V: tl.constexpr, + DK: tl.constexpr, +): + f = tl.program_id(0) + h = tl.program_id(1) + pid_v = tl.program_id(2) + rows = pid_v * BLOCK_V + tl.arange(0, BLOCK_V) + cols = tl.arange(0, DK) + valid = rows < DV + fb = tl.load(BEFORE + f) + fa = tl.load(AFTER + f) + has_b = tl.load(HASB + f) + has_a = tl.load(HASA + f) + plane = rows[:, None] * DK + cols[None, :] + off_b = ((fb * H + h) * DV) * DK + plane + off_a = ((fa * H + h) * DV) * DK + plane + sb = tl.load(PREFIX + off_b, mask=valid[:, None], other=0.0) + sa = tl.load(SUFFIX + off_a, mask=valid[:, None], other=0.0) + if HAS_TEXT: + ts = tl.load(TEXT + (h * DV) * DK + plane, mask=valid[:, None], other=0.0) + sb = tl.where(has_b != 0, sb, ts) + sa = tl.where(has_a != 0, sa, ts) + else: + sb = tl.where(has_b != 0, sb, 0.0) + sa = tl.where(has_a != 0, sa, 0.0) + if BRIDGE: + bb = tl.load(BRIDGEB + f) + ba = tl.load(BRIDGEA + f) + lp_t1 = tl.load(LOGP + ((f + 1) * H + h) * DK + cols) + lp_t = tl.load(LOGP + (f * H + h) * DK + cols) + lp_bb = tl.load(LOGP + (bb * H + h) * DK + cols) + lp_ba = tl.load(LOGP + (ba * H + h) * DK + cols) + sb = sb * tl.exp(lp_t1 - lp_bb)[None, :] + sa = sa * tl.exp(lp_ba - lp_t)[None, :] + out = sb + sa + tl.store( + OUT + ((f * H + h) * DV) * DK + plane, + out.to(OUT.dtype.element_ty), + mask=valid[:, None], + ) + + +def vdn_gather_linear_state( + prefix: torch.Tensor, + suffix: torch.Tensor, + alpha: torch.Tensor, + text_state: torch.Tensor | None, + *, + before_idx: torch.Tensor, + after_idx: torch.Tensor, + has_before: torch.Tensor, + has_after: torch.Tensor, + bridge_before: torch.Tensor, + bridge_after: torch.Tensor, + bridge: bool, + out_dtype: torch.dtype, +) -> torch.Tensor: + """The boundary gather of the linear branch as one kernel over the + [F, H, dv, dk] fp32 state banks.""" + if not prefix.is_cuda: + raise ValueError( + "vdn_gather_linear_state is a Triton kernel; inputs must be on CUDA" + ) + F, H, DV, DK = prefix.shape + _check_head_dim(DK) + prefix = prefix.contiguous() + suffix = suffix.contiguous() + if bridge: + log_alpha = torch.log(alpha.float().clamp_min(1e-12)) + logp = torch.cat( + [torch.zeros_like(log_alpha[:1]), log_alpha.cumsum(0)] + ).contiguous() + else: + logp = prefix # unused + text = text_state.float().contiguous() if text_state is not None else prefix + out = torch.empty(prefix.shape, dtype=out_dtype, device=prefix.device) + _gather_state_kernel[(F, H, triton.cdiv(DV, _BLOCK_ROWS))]( + prefix, + suffix, + logp, + text, + _i32(before_idx), + _i32(after_idx), + _i32(has_before), + _i32(has_after), + _i32(bridge_before), + _i32(bridge_after), + out, + H, + DV, + HAS_TEXT=text_state is not None, + BRIDGE=bridge, + BLOCK_V=_BLOCK_ROWS, + DK=DK, + num_warps=4, + ) + return out + + +__all__ = [ + "can_use_vdn_frame_stats_prep", + "can_use_vdn_gather_linear_state", + "can_use_vdn_linear_epilogue", + "can_use_vdn_silu_l2norm", + "can_use_vdn_temporal_conv_act", + "vdn_frame_stats_prep", + "vdn_gather_linear_state", + "vdn_linear_epilogue", + "vdn_silu_l2norm", + "vdn_temporal_conv_act", +] diff --git a/python/sglang/kernels/ops/diffusion/quantization/__init__.py b/python/sglang/kernels/ops/diffusion/quantization/__init__.py new file mode 100644 index 000000000..d3b785876 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/quantization/__init__.py @@ -0,0 +1 @@ +"""Block-scaled (MXFP8) activation quantizers with GEMM-ready scale layouts.""" diff --git a/python/sglang/kernels/ops/diffusion/quantization/mxfp8_swizzled_triton.py b/python/sglang/kernels/ops/diffusion/quantization/mxfp8_swizzled_triton.py new file mode 100644 index 000000000..13ad76df6 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/quantization/mxfp8_swizzled_triton.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 producers: e4m3 payload plus one E8M0 scale per 32 elements along K, +the scales in the cuBLASLt ``SWIZZLE_32_4_4`` layout that +``torch.nn.functional.scaled_mm(..., BlockWise1x32)`` consumes on SM100. + +Scale ``(r, c)`` of the ``[rows, K/32]`` scale matrix lives at byte +``((r // 128) * ceil(K/32 / 4) + c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4``, +rows padded to 128 and scale columns to 4, padding zero. Exponent +``e = ceil(log2(amax / 448))`` exactly from the float bits; ``q = e4m3(x * 2**-e)``; +scale byte ``e + 127``. Every producer quantizes the bf16-rounded value the +unfused bf16 kernel stores, so each is byte-exact against that kernel followed +by ``mxfp8_quantize_swizzled`` (itself byte-exact vs ``flashinfer.mxfp8_quantize``). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from sglang.kernels.ops.diffusion.common.numerics import round_bf16_to_fp32 + +_E4M3 = torch.float8_e4m3fn + + +def _scale_numel(rows: int, k: int) -> int: + n_groups = k // 32 + return -(-rows // 128) * 128 * (-(-n_groups // 4) * 4) + + +@triton.jit +def _mx_e8m0_from_amax(amax): + bits = amax.to(tl.int32, bitcast=True) + e0 = ((bits >> 23) & 0xFF) - 135 + # amax / 2**e0 lies in [256, 512); bump e0 when it exceeds 448 = 1.75 * 2**8 + thr = (((e0 + 135) << 23) | 0x600000).to(tl.float32, bitcast=True) + e = e0 + (amax > thr).to(tl.int32) + e = tl.maximum(e, -127) + inv = ((127 - e) << 23).to(tl.float32, bitcast=True) + return e + 127, inv + + +@triton.jit +def _mx_scale_offsets(r, c, n_col_blocks): + tile = (r // 128) * n_col_blocks + (c // 4) + return tile * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (c % 4) + + +@triton.jit +def _mxfp8_quant_kernel( + x_ptr, + q_ptr, + s_ptr, + rows, + k, + n_groups, + n_col_blocks, + stride_x, + BLOCK_R: tl.constexpr, + G: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_g = tl.program_id(1) + r = pid_r * BLOCK_R + tl.arange(0, BLOCK_R) + g = pid_g * G + tl.arange(0, G) + c = pid_g * (G * 32) + tl.arange(0, G * 32) + rmask = r < rows + mask = rmask[:, None] & (c < k)[None, :] + x = tl.load( + x_ptr + r[:, None].to(tl.int64) * stride_x + c[None, :], mask=mask, other=0.0 + ).to(tl.float32) + x3 = tl.reshape(x, [BLOCK_R, G, 32]) + amax = tl.max(tl.abs(x3), axis=2) + sbyte, inv = _mx_e8m0_from_amax(amax) + q = tl.reshape(x3 * inv[:, :, None], [BLOCK_R, G * 32]) + tl.store( + q_ptr + r[:, None].to(tl.int64) * k + c[None, :], + q.to(tl.float8e4nv), + mask=mask, + ) + smask = rmask[:, None] & (g < n_groups)[None, :] + tl.store( + s_ptr + _mx_scale_offsets(r[:, None], g[None, :], n_col_blocks), + sbyte.to(tl.uint8), + mask=smask, + ) + + +@triton.jit +def _silu_mul_mxfp8_kernel( + x_ptr, + q_ptr, + s_ptr, + rows, + hidden, + n_groups, + n_col_blocks, + stride_row, + BLOCK_R: tl.constexpr, + G: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) + r = pid_r * BLOCK_R + tl.arange(0, BLOCK_R) + g = pid_c * G + tl.arange(0, G) + c = pid_c * (G * 32) + tl.arange(0, G * 32) + rmask = r < rows + mask = rmask[:, None] & (c < hidden)[None, :] + base = x_ptr + r[:, None].to(tl.int64) * stride_row + gate = tl.load(base + c[None, :], mask=mask, other=0.0).to(tl.float32) + up = tl.load(base + hidden + c[None, :], mask=mask, other=0.0).to(tl.float32) + act = (gate * tl.sigmoid(gate)).to(tl.bfloat16).to(tl.float32) + prod = (act * up).to(tl.bfloat16).to(tl.float32) + p3 = tl.reshape(prod, [BLOCK_R, G, 32]) + amax = tl.max(tl.abs(p3), axis=2) + sbyte, inv = _mx_e8m0_from_amax(amax) + q = tl.reshape(p3 * inv[:, :, None], [BLOCK_R, G * 32]) + tl.store( + q_ptr + r[:, None].to(tl.int64) * hidden + c[None, :], + q.to(tl.float8e4nv), + mask=mask, + ) + smask = rmask[:, None] & (g < n_groups)[None, :] + tl.store( + s_ptr + _mx_scale_offsets(r[:, None], g[None, :], n_col_blocks), + sbyte.to(tl.uint8), + mask=smask, + ) + + +@triton.jit +def _indexed_scale_shift_mxfp8_kernel( + x_ptr, + q_ptr, + s_ptr, + shift_ptr, + scale_ptr, + indices_ptr, + hidden_size, + n_groups, + n_col_blocks, + stride_x_row, + stride_shift_row, + stride_scale_row, + stride_indices, + STORE_BF16: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, BLOCK_N) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + xrow = x_ptr + row.to(tl.int64) * stride_x_row + x = tl.load(xrow + columns, mask=mask, other=0.0).to(tl.float32) + shift = tl.load( + shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0 + ).to(tl.float32) + scale = tl.load( + scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0 + ).to(tl.float32) + # the rounding points of _indexed_scale_shift_bf16_kernel + one_plus_scale = round_bf16_to_fp32(1.0 + scale) + scaled = round_bf16_to_fp32(x * one_plus_scale) + out = round_bf16_to_fp32(scaled + shift) + if STORE_BF16: + tl.store(xrow + columns, out, mask=mask) + v3 = tl.reshape(out, [BLOCK_N // 32, 32]) + amax = tl.max(tl.abs(v3), axis=1) + sbyte, inv = _mx_e8m0_from_amax(amax) + q = tl.reshape(v3 * inv[:, None], [BLOCK_N]) + tl.store( + q_ptr + row.to(tl.int64) * hidden_size + columns, + q.to(tl.float8e4nv), + mask=mask, + ) + g = tl.arange(0, BLOCK_N // 32) + tl.store( + s_ptr + _mx_scale_offsets(row, g, n_col_blocks), + sbyte.to(tl.uint8), + mask=g < n_groups, + ) + + +def can_use_mxfp8_swizzled(x: torch.Tensor) -> bool: + """Row-major bf16 CUDA 2D tensor with K % 32 == 0, outside torch.compile.""" + return ( + x.is_cuda + and x.ndim == 2 + and x.dtype == torch.bfloat16 + and x.stride(-1) == 1 + and x.shape[-1] % 32 == 0 + and not torch.compiler.is_compiling() + ) + + +def _alloc( + rows: int, k: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + q = torch.empty(rows, k, dtype=_E4M3, device=device) + s = torch.zeros(_scale_numel(rows, k), dtype=torch.uint8, device=device) + return q, s + + +def mxfp8_quantize_swizzled(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """bf16 ``[rows, k]`` -> ``(fp8 [rows, k], swizzled e8m0 scale bytes)``.""" + if not can_use_mxfp8_swizzled(x): + raise ValueError("expected a row-major bf16 CUDA [rows, k] tensor, k % 32 == 0") + rows, k = x.shape + q, s = _alloc(rows, k, x.device) + if rows == 0: + return q, s + n_groups = k // 32 + n_col_blocks = -(-n_groups // 4) + block_r, g = 32, 8 + grid = (triton.cdiv(rows, block_r), triton.cdiv(n_groups, g)) + with torch.get_device_module().device(x.device): + _mxfp8_quant_kernel[grid]( + x, + q, + s, + rows, + k, + n_groups, + n_col_blocks, + x.stride(0), + BLOCK_R=block_r, + G=g, + num_warps=4, + ) + return q, s + + +def can_use_silu_mul_mxfp8(hidden: torch.Tensor) -> bool: + return can_use_mxfp8_swizzled(hidden) and (hidden.shape[-1] // 2) % 32 == 0 + + +def silu_mul_mxfp8(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``hidden [rows, 2n]`` bf16 (gate | up) -> quantized ``silu(gate) * up``.""" + if not can_use_silu_mul_mxfp8(hidden): + raise ValueError( + "expected a row-major bf16 CUDA [rows, 2 * n] tensor, n % 32 == 0" + ) + rows, twice = hidden.shape + n = twice // 2 + q, s = _alloc(rows, n, hidden.device) + if rows == 0: + return q, s + n_groups = n // 32 + n_col_blocks = -(-n_groups // 4) + block_r, g = 16, 8 + grid = (triton.cdiv(rows, block_r), triton.cdiv(n_groups, g)) + with torch.get_device_module().device(hidden.device): + _silu_mul_mxfp8_kernel[grid]( + hidden, + q, + s, + rows, + n, + n_groups, + n_col_blocks, + hidden.stride(0), + BLOCK_R=block_r, + G=g, + num_warps=4, + ) + return q, s + + +def indexed_scale_shift_mxfp8_( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + *, + keep_bf16: bool, +) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor]: + """``x * (1 + scale[idx]) + shift[idx]`` -> ``(x | None, fp8, scales)``; with + ``keep_bf16`` the bf16 result is also written into ``x`` and returned.""" + if not can_use_mxfp8_swizzled(x): + raise ValueError( + "expected a row-major bf16 CUDA [rows, hidden] tensor, hidden % 32 == 0" + ) + rows, hidden_size = x.shape + q, s = _alloc(rows, hidden_size, x.device) + if rows == 0: + return (x if keep_bf16 else None), q, s + n_groups = hidden_size // 32 + n_col_blocks = -(-n_groups // 4) + block_n = triton.next_power_of_2(hidden_size) + with torch.get_device_module().device(x.device): + _indexed_scale_shift_mxfp8_kernel[(rows,)]( + x, + q, + s, + shift, + scale, + indices, + hidden_size, + n_groups, + n_col_blocks, + x.stride(0), + shift.stride(0), + scale.stride(0), + indices.stride(0), + STORE_BF16=keep_bf16, + BLOCK_N=block_n, + num_warps=8, + ) + return (x if keep_bf16 else None), q, s + + +__all__ = [ + "can_use_mxfp8_swizzled", + "can_use_silu_mul_mxfp8", + "indexed_scale_shift_mxfp8_", + "mxfp8_quantize_swizzled", + "silu_mul_mxfp8", +] diff --git a/python/sglang/kernels/ops/diffusion/rope/qknorm_rope_jit.py b/python/sglang/kernels/ops/diffusion/rope/qknorm_rope_jit.py index 7cf56d34b..9ba926be9 100644 --- a/python/sglang/kernels/ops/diffusion/rope/qknorm_rope_jit.py +++ b/python/sglang/kernels/ops/diffusion/rope/qknorm_rope_jit.py @@ -33,6 +33,7 @@ def _jit_qknorm_rope_module( round_norm_before_rope: bool, pack_kv: bool = False, cache_has_full_width: bool = False, + out_of_place: bool = False, ) -> Module: args = make_cpp_args( head_dim, @@ -44,8 +45,12 @@ def _jit_qknorm_rope_module( round_norm_before_rope, cache_has_full_width, ) - op_name = "qknorm_rope_pack_kv" if pack_kv else "qknorm_rope" - kernel_name = "QKNormRopePackKVKernel" if pack_kv else "QKNormRopeKernel" + if pack_kv: + op_name, kernel_name = "qknorm_rope_pack_kv", "QKNormRopePackKVKernel" + elif out_of_place: + op_name, kernel_name = "qknorm_rope_out_of_place", "QKNormRopeOutOfPlaceKernel" + else: + op_name, kernel_name = "qknorm_rope", "QKNormRopeKernel" return load_jit( op_name, *args, @@ -182,6 +187,46 @@ def fused_inplace_qknorm_rope( module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps) +@register_custom_op(mutates_args=["q_out", "k_out"]) +def fused_qknorm_rope_out_of_place( + q: torch.Tensor, + k: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + *, + is_neox: bool, + eps: float = 1e-6, + head_dim: int = 0, + rope_dim: int = 0, + round_norm_before_rope: bool = False, + cache_has_full_width: bool = False, +) -> None: + """QK-norm + RoPE from ``q``/``k`` (any strides) into ``q_out``/``k_out``; + the inputs are left untouched. Same arithmetic as the in-place kernel.""" + head_dim = head_dim or q.size(-1) + if not rope_dim: + cache_width = cos_sin_cache.size(-1) + rope_dim = cache_width // 2 if cache_has_full_width else cache_width + module = _jit_qknorm_rope_module( + head_dim, + rope_dim, + is_neox, + q.dtype, + cos_sin_cache.dtype, + round_norm_before_rope, + False, + cache_has_full_width, + True, + ) + module.qknorm_rope_out_of_place( + q, k, q_out, k_out, q_weight, k_weight, cos_sin_cache, positions, eps + ) + + @register_custom_op(mutates_args=["q", "packed_kv"]) def fused_qknorm_rope_pack_kv( q: torch.Tensor, diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py index 36527369a..322001b24 100755 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py @@ -452,6 +452,33 @@ MODELS = { ], "force_eager": True, }, + # OpenVDN paper workload: 1344x768, 14.375 s (latent_t 102), t2va, 9 grid points = 8 NFE + "vdn-h3": { + "path": "OpenVDN/vdn-minimax-h3", + "prompt": ( + "A curious raccoon peers through a vibrant field of yellow " + "sunflowers, its eyes wide with interest." + ), + "seed": 1000, + "config_overrides": { + "task": "t2va", + "conditions": [], + "target": { + "short_edge": 768, + "aspect_ratio": "16:9", + "duration_seconds": 14.375, + }, + "num_inference_steps": 9, + }, + "extra_args": [ + "--num-gpus=8", + "--quantization=fp8", + "--performance-mode=speed", + "--enable-torch-compile=false", + "--warmup-steps=2", + ], + "force_eager": True, + }, # Source-tracked extras from current registry / GPU test coverage. "longcat-image": { "path": "meituan-longcat/LongCat-Image", diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 8f4f434de..23376f06d 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/python/sglang/multimodal_gen/README.md @@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus ## Key Features SGLang Diffusion has the following features: - - Broad model support: Wan, FastWan, FLUX, Qwen-Image, LongCat-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3/LTX-2.5, MiniMax-H3, FastH3, LingBot Video MoE, LingBot World, SANA-Video/SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more + - Broad model support: Wan, FastWan, FLUX, Qwen-Image, LongCat-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3/LTX-2.5, MiniMax-H3, FastH3, VDN-H3, LingBot Video MoE, LingBot World, SANA-Video/SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more - Fast inference speed: empowered by optimized `sgl-kernel` kernels, scheduler/runtime improvements, caching acceleration, and native diffusion hot-path optimizations - Ease of use: OpenAI-compatible api, CLI, and python sdk support - Multi-platform support: diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py index 2457fd25b..efd5484df 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py @@ -2,6 +2,9 @@ from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig +from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import ( + VDNHybridAttentionArchConfig, +) MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64 MINIMAX_H3_ADALN_MODALITY_NUM = 3 @@ -48,6 +51,8 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): ), r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"blocks.\1.attn.out_proj.\2", r"^transformer_blocks\.(\d+)\.attn\.to_gate_compress\.(.*)$": r"blocks.\1.attn.to_gate_compress.\2", + # VDN-H3 hybrid attention module (see minimax_h3_vdn_attention) + r"^transformer_blocks\.(\d+)\.attn\.(linear_attention|softmax_gate|to_out_linear)\.(.*)$": r"blocks.\1.attn.hybrid.\2.\3", r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"blocks.\1.attn.q_norm.\2", r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"blocks.\1.attn.k_norm.\2", r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"blocks.\1.mlp.fc1.\2", @@ -102,6 +107,8 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): checkpoint_uses_diffusers_layout: bool = False adaln_affine_input_dim: int | None = None has_gate_compress: bool = False + # VDN-H3: None for the dense model; set from transformer/config.json + hybrid_attention: VDNHybridAttentionArchConfig | None = None def __post_init__(self) -> None: super().__post_init__() @@ -110,6 +117,10 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): if len(self.patch_size) != 3: raise ValueError(f"patch_size must have 3 values, got {self.patch_size}.") self.num_channels_latents = self.latents_dim + if isinstance(self.hybrid_attention, dict): + self.hybrid_attention = VDNHybridAttentionArchConfig.from_transform_config( + self.hybrid_attention + ) @dataclass @@ -133,6 +144,11 @@ class MiniMaxH3DiTConfig(DiTConfig): model_dict["adaln_affine_input_dim"] = source_model_dict["time_embed_dim"] model_dict["time_embed_dim"] = source_model_dict["adaln_rank"] model_dict["adaln_curve_grid"] = source_model_dict["time_table_size"] + hybrid = model_dict.get("hybrid_attention") + if isinstance(hybrid, dict): + model_dict["hybrid_attention"] = ( + VDNHybridAttentionArchConfig.from_transform_config(hybrid) + ) super().update_model_arch(model_dict) diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3_vdn.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3_vdn.py new file mode 100644 index 000000000..5f3cf8902 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3_vdn.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 hybrid attention architecture config (window softmax + linear branch).""" + +from typing import Any + +import msgspec + +VDN_H3_DELTA_RULES = ("vdn_solve", "sana_scaled", "vdn_scaled") +VDN_H3_BRIDGE_MODES = ("alpha", "none") +VDN_H3_ANCHOR_FRAME_MODES = ("none", "columns", "rows", "both") +VDN_H3_SHORT_CONV_TARGETS = ("q", "k", "v") + + +class VDNHybridAttentionArchConfig(msgspec.Struct): + """VDN-H3 hybrid attention (window softmax + frame-wise linear branch); the + resolved ``hybrid_attention`` transform config the overlay copies into + ``transformer/config.json``. A dense checkpoint has none.""" + + # frame t is in chunk t // chunk and attends chunks [c - radius, c + radius]; + # chunk 0 means a centered frame window + chunk: int = 5 + radius: int = 1 + # "both": frames 0 and F-1 dense as rows and columns, so the branch skips them + anchor_frames: str = "both" + enable_softmax_gate: bool = True + delta_rule: str = "vdn_solve" + linear_head_dim: int = 128 + bridge: str = "alpha" + a_fp32: bool = True + enable_text_state: bool = True + short_conv: tuple[str, ...] = ("k", "v") + + def __post_init__(self) -> None: + if self.delta_rule not in VDN_H3_DELTA_RULES: + raise ValueError( + f"hybrid_attention.delta_rule={self.delta_rule!r}; expected one of " + f"{VDN_H3_DELTA_RULES}" + ) + if self.bridge not in VDN_H3_BRIDGE_MODES: + raise ValueError( + f"hybrid_attention.bridge={self.bridge!r}; expected one of " + f"{VDN_H3_BRIDGE_MODES}" + ) + if self.anchor_frames not in VDN_H3_ANCHOR_FRAME_MODES: + raise ValueError( + f"hybrid_attention.anchor_frames={self.anchor_frames!r}; expected " + f"one of {VDN_H3_ANCHOR_FRAME_MODES}" + ) + if any(t not in VDN_H3_SHORT_CONV_TARGETS for t in self.short_conv) or len( + set(self.short_conv) + ) != len(self.short_conv): + raise ValueError( + f"hybrid_attention.short_conv={self.short_conv!r}; expected a " + f"distinct subset of {VDN_H3_SHORT_CONV_TARGETS}" + ) + if self.chunk < 0 or self.radius < 0: + raise ValueError("hybrid_attention.chunk and radius must be >= 0") + if self.linear_head_dim <= 0: + raise ValueError("hybrid_attention.linear_head_dim must be positive") + + @classmethod + def from_transform_config( + cls, config: dict[str, Any] + ) -> "VDNHybridAttentionArchConfig": + """Build from VDN's nested v2 transform config.""" + soft = dict(config.get("softmax_attention", {})) + lin = dict(config.get("linear_attention", {})) + short_conv = lin.get("short_conv", {"targets": []}) + targets = ( + short_conv.get("targets", []) + if isinstance(short_conv, dict) + else list(short_conv or []) + ) + return cls( + chunk=int(soft.get("chunk", 0)), + radius=int(soft["radius"]), + anchor_frames=str(config.get("anchor_frames", "none")), + enable_softmax_gate=bool(config.get("enable_softmax_gate", True)), + delta_rule=str(lin.get("delta_rule", "vdn_solve")), + linear_head_dim=int(lin["linear_head_dim"]), + bridge=str(lin.get("bridge", "alpha")), + a_fp32=bool(lin.get("a_fp32", True)), + enable_text_state=bool(lin.get("enable_text_state", False)), + short_conv=tuple(targets), + ) + + def window_bounds(self, num_frames: int) -> list[tuple[int, int]]: + """Per-frame inclusive softmax-window bounds [lo, hi], unclamped.""" + if self.chunk <= 0: + return [(t - self.radius, t + self.radius) for t in range(num_frames)] + return [ + ( + ((t // self.chunk) - self.radius) * self.chunk, + ((t // self.chunk) + self.radius + 1) * self.chunk - 1, + ) + for t in range(num_frames) + ] + + def full_cover(self, num_frames: int) -> bool: + """True when every frame's window already spans the whole clip, i.e. + the softmax branch IS dense attention and the linear branch is off.""" + return all( + lo <= 0 and hi >= num_frames - 1 + for lo, hi in self.window_bounds(num_frames) + ) + + +__all__ = ["VDNHybridAttentionArchConfig"] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3_vdn.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3_vdn.py new file mode 100644 index 000000000..4267a92bd --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3_vdn.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 pipeline config: the MiniMax-H3 deployment envelope for the hybrid +attention checkpoint.""" + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( + MiniMaxH3PipelineConfig, +) +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, + current_platform, +) + + +@dataclass +class VDNH3PipelineConfig(MiniMaxH3PipelineConfig): + """VDN-H3 (hybrid window-softmax + Video Delta linear attention, 8-NFE DMD2 + distill; t2va and fl2va on one fl2va partition): the deployment envelope; + the arch config comes from the materialized ``transformer/config.json``.""" + + def validate_quality_deployment(self, server_args) -> None: + raise ValueError( + 'quality="high" is audited only for the base MiniMax-H3 50-step ' + "4xH200 deployment; the VDN-H3 8-step hybrid checkpoint has no " + 'audited high-quality deployment. Use quality="lossless".' + ) + + def validate_server_args(self, server_args) -> None: + if server_args.model_variant is not None: + raise ValueError( + "VDN-H3 ships one weight partition (fl2va, serving t2va and " + "fl2va); --model-variant does not apply. Ref2VA was not trained; " + "use MiniMaxAI/MiniMax-H3 --model-variant ref2va." + ) + quantization = (server_args.quantization or "").lower() + if quantization in ("none", "bf16"): + server_args.quantization = None + elif ( + current_platform.is_blackwell() or current_platform.is_sm120() + ) and quantization in ("", "fp8"): + # the block-scaled GEMM exists on SM100+ only; before that fp8 stays per-channel + server_args.quantization = "mxfp8" + # an unset backend would resolve to the platform default (dense FA) + if server_args.attention_backend is None and not ( + server_args.component_attention_backends or {} + ).get("transformer"): + server_args.attention_backend = "hybrid_window_attn_h3" + selected_backend = self.resolve_transformer_attention_backend(server_args) + if selected_backend is not AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3: + # the base-H3+LoRA equivalence smoke runs plain attention on purpose + config = server_args.attention_backend_config or {} + if not bool(config.get("vdn_h3_dense_smoke", False)): + raise ValueError( + "VDN-H3 requires --attention-backend hybrid_window_attn_h3 for " + f"the transformer (got {selected_backend}); a dense backend " + "would skip the linear branch and the softmax gates and " + "produce the wrong model, not a slower one. Pass " + "--attention-backend-config '{\"vdn_h3_dense_smoke\": true}' " + "only for the base-H3+LoRA equivalence smoke." + ) + if int(server_args.ring_degree or 1) > 1: + raise ValueError( + "VDN-H3 does not support --ring-degree > 1; use Ulysses sequence " + "parallelism." + ) + if server_args.enable_torch_compile or server_args.enable_breakable_cuda_graph: + # BCG keeps one pool per captured segment and exhausts 183 GB at 104k rows + raise ValueError( + "VDN-H3 hybrid attention is not validated under torch.compile or " + "the breakable CUDA graph yet; disable them." + ) + super().validate_server_args(server_args) + + +__all__ = ["VDNH3PipelineConfig"] diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3_vdn.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3_vdn.py new file mode 100644 index 000000000..b54d646f3 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3_vdn.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 sampling params: the 8-NFE grid on the MiniMax-H3 request surface.""" + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams + + +@dataclass +class VDNH3SamplingParams(MiniMaxH3SamplingParams): + """VDN-H3: nine sigma grid points, i.e. the eight distilled DiT forwards + (VDN counts NFEs; SGLang counts sigma grid points). The turbo adapter is + only valid at 8 NFE with video shift 12 / audio shift 3 (the defaults).""" + + num_inference_steps: int = 9 + + def _validate(self) -> None: + super()._validate() + if self.num_inference_steps != 9: + raise ValueError( + "VDN-H3 is distilled for exactly nine sigma grid points (eight DiT " + f"forwards); got num_inference_steps={self.num_inference_steps}. " + "Use MiniMaxAI/MiniMax-H3 for other schedules." + ) + if self.task is not None and self.task.strip().lower() not in ( + "t2va", + "fl2va", + ): + raise ValueError( + "VDN-H3 serves t2va and fl2va; ref2va was not trained (got " + f"task={self.task!r}). Use MiniMaxAI/MiniMax-H3 --model-variant " + "ref2va for that task." + ) + + +__all__ = ["VDNH3SamplingParams"] diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 2224d5a63..f9d05c527 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -84,6 +84,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX23PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2_5 import LTX25PipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3_vdn import ( + VDNH3PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.mova import ( MOVA360PConfig, MOVA720PConfig, @@ -170,6 +173,7 @@ from sglang.multimodal_gen.configs.sample.minimax_h3 import ( FastH3SamplingParams, MiniMaxH3SamplingParams, ) +from sglang.multimodal_gen.configs.sample.minimax_h3_vdn import VDNH3SamplingParams from sglang.multimodal_gen.configs.sample.mova import ( MOVA_360P_SamplingParams, MOVA_720P_SamplingParams, @@ -350,6 +354,7 @@ KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: Dict[str, str] = { "minimaxai/minimax-h3": "MiniMaxH3Pipeline", "minimax/minimax-h3": "MiniMaxH3Pipeline", "fastvideo/fastvideo-fasth3-4-step-preview-v1-vsa-datafree": "FastH3Pipeline", + "openvdn/vdn-minimax-h3": "VDNH3Pipeline", "lerobot/pi05": "Pi05Pipeline", "pi05": "Pi05Pipeline", "pi0.5": "Pi05Pipeline", @@ -1016,6 +1021,7 @@ def _register_configs(): model_detectors=[ lambda model_id: ( "minimaxh3" in model_id.lower().replace("-", "").replace("_", "") + and "vdn" not in model_id.lower() ) ], ) @@ -1038,6 +1044,19 @@ def _register_configs(): ) ], ) + register_configs( + sampling_param_cls=VDNH3SamplingParams, + pipeline_config_cls=VDNH3PipelineConfig, + hf_model_paths=[ + "OpenVDN/vdn-minimax-h3", + ], + model_detectors=[ + lambda model_id: ( + "vdn" in model_id.lower() + and "minimaxh3" in model_id.lower().replace("-", "").replace("_", "") + ) + ], + ) # FLUX register_configs( sampling_param_cls=FluxSamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/hybrid_window_attn_h3.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/hybrid_window_attn_h3.py new file mode 100644 index 000000000..763e2eb06 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/hybrid_window_attn_h3.py @@ -0,0 +1,521 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 window-softmax backend on the MiniMax-H3 packed layout. + +An exact softmax over a chunk-aligned frame window: frame t belongs to chunk +t // chunk and attends to chunks [c - radius, c + radius]; frames 0 and F-1 +are dense anchors; text and audio rows are dense both ways; padding rows sit +outside every mask; a per-(token, head) sigmoid gate scales the output. The +linear branch (``minimax_h3_vdn.py``) covers the window's complement. The +metadata is request-static and installed once per request through +``set_forward_context``. The window runs as a union of dense varlen +FlashAttention calls: the dense-query rows against all keys, then per-chunk +gathered [globals | window | anchors] K/V; same math as a masked kernel up to +bf16 reduction order. +""" + +from __future__ import annotations + +import functools +import re +from dataclasses import dataclass +from typing import Any + +import msgspec +import torch + +from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func +from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import ( + VDNHybridAttentionArchConfig, +) +from sglang.multimodal_gen.runtime.layers.attention.backends import ( + flash_attn as _flash_attn_backend, +) +from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import VDNH3Layout +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + +_DIT_BLOCK_PREFIX = re.compile(r"^blocks\.(\d+)\.") + + +class HybridWindowAttentionH3Backend(AttentionBackend): + accept_output_buffer: bool = False + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] + + @staticmethod + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3 + + @staticmethod + def get_impl_cls() -> type[HybridWindowAttentionH3Impl]: + return HybridWindowAttentionH3Impl + + @staticmethod + def get_metadata_cls() -> type[HybridWindowAttentionH3Metadata]: + return HybridWindowAttentionH3Metadata + + @staticmethod + def get_builder_cls() -> type[HybridWindowAttentionH3MetadataBuilder]: + return HybridWindowAttentionH3MetadataBuilder + + +def window_mask_frames( + hybrid: VDNHybridAttentionArchConfig, num_frames: int +) -> tuple[list[tuple[int, int]], set[int], set[int]]: + """(clamped per-frame window bounds, dense-ROW frames, dense-COLUMN frames).""" + bounds = [ + (max(lo, 0), min(hi, num_frames - 1)) + for lo, hi in hybrid.window_bounds(num_frames) + ] + anchors = {0, num_frames - 1} if hybrid.anchor_frames != "none" else set() + dense_rows = anchors if hybrid.anchor_frames in ("rows", "both") else set() + dense_cols = anchors if hybrid.anchor_frames in ("columns", "both") else set() + return bounds, dense_rows, dense_cols + + +def window_mask_reference( + hybrid: VDNHybridAttentionArchConfig, layout: VDNH3Layout, device: torch.device +) -> torch.Tensor: + """Dense boolean [used, used] mask of the softmax branch, for tests.""" + used = layout.used + keep = torch.ones(used, used, dtype=torch.bool, device=device) + vs, ve = layout.video_start, layout.video_end + bounds, dense_rows, dense_cols = window_mask_frames(hybrid, layout.num_frames) + tpf = layout.tokens_per_frame + rows = torch.arange(vs, ve, device=device) + frame_of = (rows - vs) // tpf + qf = frame_of[:, None] + kf = frame_of[None, :] + lo = torch.tensor([b[0] for b in bounds], device=device)[qf] + hi = torch.tensor([b[1] for b in bounds], device=device)[qf] + inside = (kf >= lo) & (kf <= hi) + for f in dense_rows: + inside |= qf == f + for f in dense_cols: + inside |= kf == f + keep[vs:ve, vs:ve] = inside + return keep + + +def _merge_ranges(ranges: list[tuple[int, int]]) -> list[tuple[int, int]]: + out: list[tuple[int, int]] = [] + for a, b in sorted(ranges): + if out and out[-1][1] >= a: + out[-1] = (out[-1][0], max(out[-1][1], b)) + else: + out.append((a, b)) + return out + + +def _cat_ranges(ranges: list[tuple[int, int]], *, device: torch.device) -> torch.Tensor: + if not ranges: + return torch.empty(0, dtype=torch.long, device=device) + return torch.cat( + [torch.arange(a, b, device=device, dtype=torch.long) for a, b in ranges] + ) + + +def _chunk_groups( + raw_bounds: list[tuple[int, int]], dense_rows: set[int] +) -> list[list[int]]: + # consecutive window frames with identical bounds share one varlen segment + groups: list[list[int]] = [] + for f in range(len(raw_bounds)): + if f in dense_rows: + continue + if ( + groups + and raw_bounds[groups[-1][-1]] == raw_bounds[f] + and groups[-1][-1] == f - 1 + ): + groups[-1].append(f) + else: + groups.append([f]) + return groups + + +class _ChunkGroup(msgspec.Struct, frozen=True): + frames: list[int] + query_rows: torch.Tensor + kv_rows: torch.Tensor + + +class _WindowPass(msgspec.Struct, frozen=True): + query_rows: torch.Tensor + query_slice: tuple[int, int] | None # set when the query rows are contiguous + kv_rows: torch.Tensor + cu_q: torch.Tensor + cu_k: torch.Tensor + max_q: int + max_k: int + + +def _window_pass( + layout: VDNH3Layout, groups: list[_ChunkGroup], device: torch.device +) -> _WindowPass: + query_lens = [int(group.query_rows.numel()) for group in groups] + kv_lens = [int(group.kv_rows.numel()) for group in groups] + frames = [frame for group in groups for frame in group.frames] + contiguous = frames == list(range(frames[0], frames[0] + len(frames))) + zero = torch.zeros(1, dtype=torch.long) + return _WindowPass( + query_rows=torch.cat([group.query_rows for group in groups]), + query_slice=( + (layout.frame_rows(frames[0])[0], layout.frame_rows(frames[-1])[1]) + if contiguous + else None + ), + kv_rows=torch.cat([group.kv_rows for group in groups]), + cu_q=torch.cat([zero, torch.tensor(query_lens).cumsum(0)]).to( + device, torch.int32 + ), + cu_k=torch.cat([zero, torch.tensor(kv_lens).cumsum(0)]).to(device, torch.int32), + max_q=max(query_lens), + max_k=max(kv_lens), + ) + + +def _window_passes( + layout: VDNH3Layout, + groups: list[_ChunkGroup], + max_gather_rows: int, + device: torch.device, +) -> list[_WindowPass]: + # one varlen call per pass; consecutive chunk groups fill up to max_gather_rows + passes: list[_WindowPass] = [] + current: list[_ChunkGroup] = [] + current_rows = 0 + for group in groups: + rows = int(group.kv_rows.numel()) + if current and current_rows + rows > max_gather_rows: + passes.append(_window_pass(layout, current, device)) + current, current_rows = [], 0 + current.append(group) + current_rows += rows + if current: + passes.append(_window_pass(layout, current, device)) + return passes + + +class _DecomposedPlan: + """Query-row groups with identical kept key sets, as dense varlen calls: + the dense-q rows against all ``used`` keys, then each chunk of frames + against its gathered [globals | window | anchors] keys.""" + + __slots__ = ("dense_q", "dense_cu_q", "dense_cu_k", "passes") + + def __init__( + self, + layout: VDNH3Layout, + hybrid: VDNHybridAttentionArchConfig, + device: torch.device, + max_gather_rows: int = 200_000, + ) -> None: + used, num_frames = layout.used, layout.num_frames + bounds, dense_rows, dense_cols = window_mask_frames(hybrid, num_frames) + rows = functools.partial(_cat_ranges, device=device) + dense_ranges = _merge_ranges( + layout.global_ranges + [layout.frame_rows(f) for f in sorted(dense_rows)] + ) + self.dense_q = rows(dense_ranges) + # built once: a tensor from a Python list costs a pageable H2D copy + sync + self.dense_cu_q = torch.tensor( + [0, int(self.dense_q.numel())], dtype=torch.int32, device=device + ) + self.dense_cu_k = torch.tensor([0, used], dtype=torch.int32, device=device) + groups = [] + for frames in _chunk_groups(hybrid.window_bounds(num_frames), dense_rows): + lo, hi = bounds[frames[0]] + kv_frames = sorted(set(range(lo, hi + 1)) | dense_cols) + groups.append( + _ChunkGroup( + frames=frames, + query_rows=rows( + _merge_ranges([layout.frame_rows(f) for f in frames]) + ), + kv_rows=rows( + _merge_ranges( + layout.global_ranges + + [layout.frame_rows(f) for f in kv_frames] + ) + ), + ) + ) + self.passes = _window_passes(layout, groups, max_gather_rows, device) + window_rows = sum(int(p.query_rows.numel()) for p in self.passes) + covered = int(self.dense_q.numel()) + window_rows + if covered != used: + raise ValueError( + f"window decomposition covers {covered} of {used} packed rows" + ) + + +@dataclass +class HybridWindowAttentionH3Metadata(AttentionMetadata): + layout: VDNH3Layout + # radius >= F: the window IS dense attention and the linear branch is off + full_cover: bool + decomposed: _DecomposedPlan | None = None + # (cos_sin [seq_len, rope_dim] bf16, positions [seq_len]) under Ulysses, else None + rope_cache_full: tuple[torch.Tensor, torch.Tensor] | None = None + + +class HybridWindowAttentionH3MetadataBuilder(AttentionMetadataBuilder): + def __init__(self) -> None: + pass + + def prepare(self) -> None: + pass + + def build( # type: ignore[override] + self, + *, + layout: VDNH3Layout, + hybrid: VDNHybridAttentionArchConfig, + device: torch.device, + rope_cache_full: tuple[torch.Tensor, torch.Tensor] | None = None, + current_timestep: int = 0, + max_gather_rows: int = 200_000, + **kwargs: dict[str, Any], + ) -> HybridWindowAttentionH3Metadata: + full_cover = hybrid.full_cover(layout.num_frames) + decomposed = None + if not full_cover: + decomposed = _DecomposedPlan( + layout, hybrid, device, max_gather_rows=max_gather_rows + ) + return HybridWindowAttentionH3Metadata( + current_timestep=current_timestep, + layout=layout, + full_cover=full_cover, + decomposed=decomposed, + rope_cache_full=rope_cache_full, + ) + + +def _fa_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_q: torch.Tensor, + cu_k: torch.Tensor, + max_q: int, + max_k: int, + scale: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + attn_out = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max_q, + max_seqlen_k=max_k, + softmax_scale=scale, + causal=False, + ver=_flash_attn_backend.fa_ver, + out=out, + ) + attn_out = attn_out[0] if isinstance(attn_out, tuple) else attn_out + if out is not None and attn_out.data_ptr() != out.data_ptr(): + out.copy_(attn_out) + return out + return attn_out + + +class HybridWindowAttentionH3Impl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.num_heads = num_heads + self.head_size = head_size + self.softmax_scale = softmax_scale + self.prefix = prefix + match = _DIT_BLOCK_PREFIX.match(prefix) + self.layer_idx = int(match.group(1)) if match else None + # non-DiT callers (the token refiner) resolve this backend too: dense FA + self._dense_fallback = _flash_attn_backend.FlashAttentionImpl( + num_heads=num_heads, + head_size=head_size, + causal=causal, + softmax_scale=softmax_scale, + num_kv_heads=num_kv_heads, + prefix=prefix, + ) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """Dense FlashAttention for the non-DiT layers this backend reaches.""" + return self._dense_fallback.forward(query, key, value, attn_metadata) + + def dense_varlen( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + cu_seqlens: torch.Tensor, + max_seqlen: int, + cu_seqlens_host: tuple[int, ...] | None = None, + ) -> torch.Tensor: + return self._dense_fallback.forward_varlen( + query, + key, + value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + ) + + def forward_varlen( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + cu_seqlens: torch.Tensor, + max_seqlen: int, + cu_seqlens_host: tuple[int, ...] | None = None, + attn_metadata: HybridWindowAttentionH3Metadata | None = None, + softmax_gate: torch.Tensor | None = None, + ) -> torch.Tensor: + """query/key/value: [T, H, D] packed rows (post-norm, post-RoPE) -> + [T, H, D]; ``softmax_gate`` [T, H] scales the output per (row, head). + Rows at and past ``used`` (padding) are zero.""" + if self.layer_idx is not None and attn_metadata is None: + raise RuntimeError( + "hybrid_window_attn_h3 needs per-request attention metadata " + "from the MiniMax-H3 denoising stage; none was set in the " + "forward context." + ) + if self.layer_idx is None: + return self.dense_varlen( + query, + key, + value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + ) + + meta = attn_metadata + layout = meta.layout + bounds = ( + cu_seqlens_host + if cu_seqlens_host is not None + else tuple(int(item) for item in cu_seqlens.tolist()) + ) + used = int(bounds[1]) + if used != layout.used or query.shape[0] != layout.seq_len: + raise ValueError( + f"hybrid_window_attn_h3 metadata was built for used={layout.used} " + f"of seq_len={layout.seq_len} rows, got used={used} of " + f"{query.shape[0]}. The request metadata and the packed layout " + "have diverged." + ) + + if meta.full_cover: + out = self.dense_varlen( + query, + key, + value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + ) + else: + out = self._decomposed(query, key, value, meta.decomposed, used) + + if softmax_gate is not None: + out.mul_(softmax_gate.to(out.dtype).unsqueeze(-1)) + if used < out.shape[0]: + out[used:].zero_() + return out + + def _decomposed( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + plan: _DecomposedPlan, + used: int, + ) -> torch.Tensor: + out = torch.empty_like(query) + key_used = key[:used] + value_used = value[:used] + if not key_used.is_contiguous(): + key_used = key_used.contiguous() + if not value_used.is_contiguous(): + value_used = value_used.contiguous() + if plan.dense_q.numel(): + out[plan.dense_q] = _fa_varlen( + torch.index_select(query, 0, plan.dense_q), + key_used, + value_used, + cu_q=plan.dense_cu_q, + cu_k=plan.dense_cu_k, + max_q=int(plan.dense_q.numel()), + max_k=used, + scale=self.softmax_scale, + ) + for window in plan.passes: + # index_select on contiguous copies takes the vectorized gather kernel + keys = torch.index_select(key_used, 0, window.kv_rows) + values = torch.index_select(value_used, 0, window.kv_rows) + if window.query_slice is not None: + start, stop = window.query_slice + _fa_varlen( + query[start:stop], + keys, + values, + cu_q=window.cu_q, + cu_k=window.cu_k, + max_q=window.max_q, + max_k=window.max_k, + scale=self.softmax_scale, + out=out[start:stop], + ) + else: + out[window.query_rows] = _fa_varlen( + torch.index_select(query, 0, window.query_rows), + keys, + values, + cu_q=window.cu_q, + cu_k=window.cu_k, + max_q=window.max_q, + max_k=window.max_k, + scale=self.softmax_scale, + ) + del keys, values + return out + + +__all__ = [ + "HybridWindowAttentionH3Backend", + "HybridWindowAttentionH3Impl", + "HybridWindowAttentionH3Metadata", + "HybridWindowAttentionH3MetadataBuilder", + "window_mask_frames", + "window_mask_reference", +] diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 5865b6896..ec8ae6b68 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -91,6 +91,11 @@ def adjust_scalar_to_fused_array( class LinearMethodBase(QuantizeMethodBase): """Base class for different (maybe quantized) linear methods.""" + def accepts_mxfp8_input(self, layer: torch.nn.Module) -> bool: + """Whether ``apply`` takes a prequantized ``(e4m3 input, swizzled E8M0 + block scales)`` tuple for this layer.""" + return False + @abstractmethod def create_weights( self, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py index 3a2396650..3f0feacee 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py @@ -76,6 +76,12 @@ class MXFP8Config(SRTFp8Config, QuantizationConfig): return UnquantizedLinearMethod() if current_platform.is_npu(): return NPUMXFP8LinearMethod(self) + if not self.is_checkpoint_fp8_serialized: + from sglang.multimodal_gen.runtime.layers.quantization.mxfp8_online import ( + MXFP8OnlineLinearMethod, + ) + + return MXFP8OnlineLinearMethod(self) return SRTFp8LinearMethod(self) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_online.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_online.py new file mode 100644 index 000000000..7feba6906 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_online.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Online MXFP8 for diffusion linears: ``--quantization mxfp8`` on a bf16 checkpoint. + +Weights quantize at load to e4m3 with one E8M0 scale per 32 elements along K +(scales in the cuBLASLt 128x4 swizzled layout, ``mxfp8_quantize_swizzled``); +activations take the same block quant per call unless the producer hands over +a prequantized ``(fp8, swizzled scales)`` tuple; the GEMM is cuBLASLt's +block-scaled ``torch.nn.functional.scaled_mm``. Layers with K not a multiple +of 32, N not a multiple of 16, non-bf16 params, or pre-Blackwell GPUs keep +the per-channel fp8 path. +""" + +from typing import Optional + +import torch +from torch.nn import Module +from torch.nn.functional import ScalingType, SwizzleType, scaled_mm + +from sglang.kernels.ops.diffusion import ( + can_use_mxfp8_swizzled, + mxfp8_quantize_swizzled, +) +from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8LinearMethod + +_E8M0 = torch.float8_e8m0fnu + + +def mxfp8_scaled_mm( + a: torch.Tensor, + a_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + *, + bias: Optional[torch.Tensor], + output_dtype: torch.dtype, +) -> torch.Tensor: + return scaled_mm( + a, + weight.t(), + scale_a=a_scale, + scale_recipe_a=ScalingType.BlockWise1x32, + swizzle_a=SwizzleType.SWIZZLE_32_4_4, + scale_b=weight_scale, + scale_recipe_b=ScalingType.BlockWise1x32, + swizzle_b=SwizzleType.SWIZZLE_32_4_4, + bias=bias, + output_dtype=output_dtype, + ) + + +class MXFP8OnlineLinearMethod(Fp8LinearMethod): + def __init__(self, quant_config) -> None: + super().__init__(quant_config) + # SRTFp8Config(use_mxfp8) sets a block size; the fallback path loads bf16 + self.block_quant = False + + def process_weights_after_loading(self, layer: Module) -> None: + layer.mxfp8 = ( + not self.use_marlin + and can_use_mxfp8_swizzled(layer.weight) + and torch.cuda.get_device_capability(layer.weight.device)[0] >= 10 + and layer.weight.shape[0] % 16 == 0 + ) + if not layer.mxfp8: + super().process_weights_after_loading(layer) + return + qweight, scale = mxfp8_quantize_swizzled(layer.weight.data.contiguous()) + layer.weight = torch.nn.Parameter(qweight, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(scale.view(_E8M0), requires_grad=False) + layer.input_scale = None + + def accepts_mxfp8_input(self, layer: Module) -> bool: + return bool(layer.mxfp8) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not layer.mxfp8: + return super().apply(layer, x, bias) + if isinstance(x, tuple): + a, a_scale = x + a_scale = a_scale.view(_E8M0) + lead, output_dtype = a.shape[:-1], torch.bfloat16 + else: + lead, output_dtype = x.shape[:-1], x.dtype + a, a_scale = mxfp8_quantize_swizzled( + x.reshape(-1, x.shape[-1]).contiguous() + ) + a_scale = a_scale.view(_E8M0) + out = mxfp8_scaled_mm( + a.reshape(-1, a.shape[-1]), + a_scale, + layer.weight, + layer.weight_scale, + bias=bias, + output_dtype=output_dtype, + ) + return out.view(*lead, out.shape[-1]) + + +__all__ = ["MXFP8OnlineLinearMethod", "mxfp8_scaled_mm"] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index e49aed982..db1d29a2d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -22,10 +22,14 @@ from sglang.kernels.ops.activation.activation import ( ) from sglang.kernels.ops.diffusion import ( can_use_fused_inplace_qknorm_rope, + can_use_mxfp8_swizzled, + can_use_silu_mul_mxfp8, fused_inplace_qknorm_rope, indexed_gate_bf16, indexed_gate_bf16_, indexed_scale_shift_bf16_, + indexed_scale_shift_mxfp8_, + silu_mul_mxfp8, ) from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm from sglang.multimodal_gen import envs @@ -79,6 +83,9 @@ from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import ( from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import ( native_adaln_weight_files, ) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import ( + MiniMaxH3VDNHybridAttention, +) from sglang.multimodal_gen.runtime.models.parameter import BlockQuantScaleParameter from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, @@ -362,6 +369,12 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) +def _accepts_mxfp8_input(linear: nn.Module) -> bool: + return linear.quant_method is not None and linear.quant_method.accepts_mxfp8_input( + linear + ) + + def _modulate_scale_shift( x: torch.Tensor, shift: torch.Tensor, @@ -827,6 +840,10 @@ class MiniMaxH3Attention(nn.Module): quant_config=None, prefix=f"{prefix}.to_gate_compress", ) + # VDN-H3: None for the dense model and for the token refiner + self.hybrid = MiniMaxH3VDNHybridAttention.build( + arch, quant_config, prefix=prefix, local_heads=self.num_heads + ) def _set_attention_backend(self, backend) -> None: if ( @@ -1032,9 +1049,12 @@ class MiniMaxH3Attention(nn.Module): subblock_sparse_query_block_mask: torch.Tensor | None = None, ulysses_active: bool = False, ring_active: bool = False, + x_prequant: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: """x: [T, hidden] packed thd rows -> [T, hidden]. + ``x_prequant``: ``x`` already quantized for ``qkv_proj`` as ``(fp8, scales)``. + Operation order: fused qkv projection -> per-head q/k RMSNorm -> RoPE on q/k -> variable-length non-causal flash attention -> output projection. @@ -1054,11 +1074,29 @@ class MiniMaxH3Attention(nn.Module): ) total = x.shape[0] - qkv, _ = self.qkv_proj(x) + qkv, _ = self.qkv_proj(x if x_prequant is None else x_prequant) q, k, v = qkv.split(self.local_inner_dim, dim=-1) q = q.view(total, self.num_heads, self.head_dim) k = k.view(total, self.num_heads, self.head_dim) v = v.view(total, self.num_heads, self.head_dim) + if ( + self.hybrid is not None + and self._attention_backend_enum + is AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3 + ): + return self.hybrid( + self, + x, + q, + k, + v, + rope_cache=rope_cache, + cu_seqlens=cu_seqlens, + cu_seqlens_host=cu_seqlens_host, + max_seqlen=max_seqlen, + ulysses_active=ulysses_active, + ring_active=ring_active, + ) if rope_cache is None: q, k = _apply_qk_norm( q, @@ -1160,7 +1198,7 @@ class MiniMaxH3MLP(nn.Module): ) def forward(self, x: torch.Tensor) -> torch.Tensor: - if x.device.type == "mps": + if not isinstance(x, tuple) and x.device.type == "mps": out = torch.empty_like(x) for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE): stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0]) @@ -1173,6 +1211,9 @@ class MiniMaxH3MLP(nn.Module): torch.mps.empty_cache() return out hidden, _ = self.fc1(x) + if _accepts_mxfp8_input(self.fc2) and can_use_silu_mul_mxfp8(hidden): + out, _ = self.fc2(silu_mul_mxfp8(hidden)) + return out hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation) out, _ = self.fc2(hidden) return out @@ -1384,11 +1425,20 @@ class MiniMaxH3DiTBlock(nn.Module): # a block-local buffer. residual = x h = self.norm1(x) - h = _modulate_scale_shift( - h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE - ) + h_prequant = None + if _accepts_mxfp8_input(self.attn.qkv_proj) and can_use_mxfp8_swizzled(h): + # the bf16 modulated rows stay in h for the VDN branch projections + h, h_fp8, h_scales = indexed_scale_shift_mxfp8_( + h, shift_msa, scale_msa, combined_indices, keep_bf16=True + ) + h_prequant = (h_fp8, h_scales) + else: + h = _modulate_scale_shift( + h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE + ) h = self.attn( h, + x_prequant=h_prequant, rope_cache=rope_cache, cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, @@ -1408,10 +1458,16 @@ class MiniMaxH3DiTBlock(nn.Module): residual = x h = self.norm2(x) - h = _modulate_scale_shift( - h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE - ) - h = self.mlp(h) + if _accepts_mxfp8_input(self.mlp.fc1) and can_use_mxfp8_swizzled(h): + _, h_fp8, h_scales = indexed_scale_shift_mxfp8_( + h, shift_mlp, scale_mlp, combined_indices, keep_bf16=False + ) + h = self.mlp((h_fp8, h_scales)) + else: + h = _modulate_scale_shift( + h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE + ) + h = self.mlp(h) # `residual` is block-local here (see above), so this stays in-place # even while Cache-DiT is attached. return _modulate_gate( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn.py new file mode 100644 index 000000000..4c0accace --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn.py @@ -0,0 +1,1078 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 (Video DeltaNet MiniMax-H3) linear attention branch. + +Port of OpenVDN's ``BidirectionalLinearBranch`` (github.com/OpenVDN/ +vdn-minimax-h3, ``src/models/linear_attention``) to SGLang's packed MiniMax-H3 +attention. The branch summarises everything the chunked window softmax cannot +see, for every video token, in five steps: + + 0. text state the prompt rows written once into a zero state; both + directional scans start from half of it + 1. features SiLU (+ separable 5x5 spatial / 5-tap temporal depthwise + conv on k, v), L2-normalised q/k, NoPE + 2. frame stats A = K^T diag(beta) K (fp32), B = V^T diag(beta) K per frame + 3. two scans Video Delta rule S_t = (S_{t-1} diag(alpha_t) + B_t)(I + A_t)^-1 + forward and reverse over frames + 4. boundary gather prefix[lo-1] + suffix[hi+1] decayed to frame t by + prod alpha over the window (the exact complement of the + softmax window; ends read the text state) + 5. readout q . S -> RMSNorm(head_dim) -> low-rank sigmoid gate + +Inference-only and eager. Parameter names follow VDN one level below +``blocks.N.attn.linear_attention``; heads shard under TP, the per-token +``alpha.down`` and ``output_gate.down`` are replicated. +""" + +from __future__ import annotations + +import functools +import math + +import msgspec +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.kernels.ops.diffusion import ( + can_use_vdn_delta_factors, + can_use_vdn_frame_stats_prep, + can_use_vdn_gather_linear_state, + can_use_vdn_linear_epilogue, + can_use_vdn_silu_l2norm, + can_use_vdn_temporal_conv_act, + vdn_delta_factors, + vdn_frame_stats_prep, + vdn_gather_linear_state, + vdn_linear_epilogue, + vdn_silu_l2norm, + vdn_temporal_conv_act, +) +from sglang.multimodal_gen.configs.models.dits.minimax_h3 import ( + MiniMaxH3DiTArchConfig, + VDNHybridAttentionArchConfig, +) +from sglang.multimodal_gen.runtime.distributed import ( + get_tp_rank, + get_tp_world_size, +) +from sglang.multimodal_gen.runtime.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, +) + +_BF16 = torch.bfloat16 +_FP32 = torch.float32 + +# Each directional scan starts from TEXT_STATE_SCALE * S_text. Baked into the +# trained checkpoints, not a knob (see VDN BidirectionalLinearBranch). +TEXT_STATE_SCALE = 0.5 +SHORT_CONV_KERNEL = 5 + + +# -------------------------------------------------------------------------- +# Packed-sequence geometry +# -------------------------------------------------------------------------- + + +class VDNH3Layout(msgspec.Struct, frozen=True): + """Where the modalities sit in SGLang's packed H3 sequence + ``[text L | cond C | audio A | video V | pad P]``. + + ``used`` is ``cu_seqlens[1]``: rows at and past it are padding and sit + outside every attention mask. Text, condition and audio rows are "global" + for the softmax branch (dense both ways); only the text rows seed the + linear branch's state. + """ + + seq_len: int + used: int + text_len: int + video_start: int + num_frames: int + tokens_per_frame: int + frame_height: int + frame_width: int + + def __post_init__(self) -> None: + if self.frame_height * self.frame_width != self.tokens_per_frame: + raise ValueError( + f"frame grid {self.frame_height}x{self.frame_width} != " + f"{self.tokens_per_frame} tokens per frame" + ) + if self.video_end > self.used or self.used > self.seq_len: + raise ValueError( + f"video rows [{self.video_start}, {self.video_end}) exceed used " + f"rows {self.used} (seq_len {self.seq_len})" + ) + if self.text_len > self.video_start: + raise ValueError("text rows must precede the video rows") + + @property + def video_end(self) -> int: + return self.video_start + self.num_frames * self.tokens_per_frame + + @property + def frame_size(self) -> tuple[int, int]: + return self.frame_height, self.frame_width + + @property + def global_ranges(self) -> list[tuple[int, int]]: + """Non-video, non-padding row ranges (text, condition, audio).""" + return [ + (start, stop) + for start, stop in ((0, self.video_start), (self.video_end, self.used)) + if start < stop + ] + + def frame_rows(self, frame: int) -> tuple[int, int]: + start = self.video_start + frame * self.tokens_per_frame + return start, start + self.tokens_per_frame + + +def vdn_h3_layout_from_packed( + packed: dict, *, latent_t: int, latent_h: int, latent_w: int +) -> VDNH3Layout: + """Layout from a ``minimax_h3_packed_sequence`` result (t2va / fl2va).""" + img_pos = packed["img_pos"].view(-1) + update_mask = packed["update_mask"].view(-1).to(torch.bool) + video_pos = img_pos[update_mask] + frame_h, frame_w = latent_h // 2, latent_w // 2 + tokens_per_frame = frame_h * frame_w + video_start = int(video_pos[0]) + if int(video_pos[-1]) - video_start + 1 != int(video_pos.numel()): + raise ValueError("video rows are not contiguous in the packed sequence") + if int(video_pos.numel()) != latent_t * tokens_per_frame: + raise ValueError( + f"{int(video_pos.numel())} video rows != {latent_t} frames x " + f"{tokens_per_frame} tokens per frame" + ) + cu = packed["cu_seqlens"].view(-1).tolist() + return VDNH3Layout( + seq_len=int(packed["seq_len"]), + used=int(cu[1]), + text_len=int(packed["text_pos"].numel()), + video_start=video_start, + num_frames=latent_t, + tokens_per_frame=tokens_per_frame, + frame_height=frame_h, + frame_width=frame_w, + ) + + +# -------------------------------------------------------------------------- +# Head-sharded plain parameters +# -------------------------------------------------------------------------- + + +def _head_sharded_loader(shard_dim: int): + # head-major output rows: TP shards by head, the checkpoint stores all heads + + def _loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_size = get_tp_world_size() + if tp_size > 1: + shard = param.shape[shard_dim] + loaded_weight = loaded_weight.narrow( + shard_dim, get_tp_rank() * shard, shard + ) + assert param.shape == loaded_weight.shape, ( + f"VDN branch parameter shape {tuple(param.shape)} != checkpoint " + f"{tuple(loaded_weight.shape)}" + ) + param.data.copy_(loaded_weight) + + return _loader + + +def _make_param( + shape: tuple[int, ...], *, dtype: torch.dtype, shard_dim: int | None +) -> nn.Parameter: + param = nn.Parameter(torch.empty(shape, dtype=dtype), requires_grad=False) + if shard_dim is not None: + param.weight_loader = _head_sharded_loader(shard_dim) + return param + + +def _head_sharded_linear( + in_features: int, out_features: int, *, bias: bool, prefix: str +) -> ColumnParallelLinear: + return ColumnParallelLinear( + in_features, + out_features, + bias=bias, + gather_output=False, + params_dtype=_BF16, + quant_config=None, + prefix=prefix, + ) + + +def _replicated_linear( + in_features: int, out_features: int, *, prefix: str +) -> ReplicatedLinear: + return ReplicatedLinear( + in_features, + out_features, + bias=False, + params_dtype=_BF16, + quant_config=None, + prefix=prefix, + ) + + +class VDNFrameAlpha(nn.Module): + """alpha_t = exp(-exp(A_log) * softplus(up(down(frame_mean)) + dt_bias)), + per frame / head / key channel, in fp32 (KDA's double-exponential gate).""" + + def __init__( + self, + hidden_size: int, + heads: int, + local_heads: int, + head_dim: int, + *, + prefix: str, + ) -> None: + super().__init__() + self.local_heads, self.head_dim = local_heads, head_dim + self.down = _replicated_linear(hidden_size, head_dim, prefix=f"{prefix}.down") + self.up = _head_sharded_linear( + head_dim, heads * head_dim, bias=False, prefix=f"{prefix}.up" + ) + # fp32: the scan multiplies alpha over ~100 frames, so bf16 error compounds + self.A_log = _make_param((local_heads,), dtype=_FP32, shard_dim=0) + self.dt_bias = _make_param((local_heads * head_dim,), dtype=_FP32, shard_dim=0) + + def forward( + self, frame_mean: torch.Tensor, heads: slice | None = None + ) -> torch.Tensor: + """frame_mean [F, hidden] fp32 -> alpha [F, H, d] fp32 for the head + range ``heads`` (Ulysses: this rank's shard of the TP-local heads).""" + if heads is None: + up_w, dt_bias, a_log, n_heads = ( + self.up.weight, + self.dt_bias, + self.A_log, + self.local_heads, + ) + else: + rows = slice(heads.start * self.head_dim, heads.stop * self.head_dim) + up_w, dt_bias, a_log = ( + self.up.weight[rows], + self.dt_bias[rows], + self.A_log[heads], + ) + n_heads = heads.stop - heads.start + delta = F.linear(frame_mean.float(), self.down.weight.float()) + delta = F.linear(delta, up_w.float()) + dt_bias + scale = torch.exp(a_log)[:, None] + delta = delta.view(-1, n_heads, self.head_dim) + return torch.exp(-scale * F.softplus(delta)) + + +class VDNOutputGate(nn.Module): + """Low-rank sigmoid gate: sigmoid(up(down(x))) -> [T, H_local, d].""" + + def __init__( + self, + hidden_size: int, + heads: int, + local_heads: int, + head_dim: int, + *, + prefix: str, + ) -> None: + super().__init__() + self.local_heads, self.head_dim = local_heads, head_dim + self.down = _replicated_linear(hidden_size, head_dim, prefix=f"{prefix}.down") + self.up = _head_sharded_linear( + head_dim, heads * head_dim, bias=True, prefix=f"{prefix}.up" + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + hidden, _ = self.down(x) + return self.up_gate(hidden) + + def up_gate(self, hidden: torch.Tensor, heads: slice | None = None) -> torch.Tensor: + """sigmoid(up(hidden)) -> [T, h, d]; ``heads`` selects a head range of + the up projection (Ulysses computes the gate on its head shard from the + all-gathered ``down`` hidden).""" + if heads is None: + gate, _ = self.up(hidden) + return torch.sigmoid(gate).view(-1, self.local_heads, self.head_dim) + rows = slice(heads.start * self.head_dim, heads.stop * self.head_dim) + bias = None if self.up.bias is None else self.up.bias[rows] + gate = F.linear(hidden, self.up.weight[rows], bias) + return torch.sigmoid(gate).view(-1, heads.stop - heads.start, self.head_dim) + + +class VDNSoftmaxGate(nn.Module): + """Per-(token, head) sigmoid gate on the softmax branch output.""" + + def __init__(self, hidden_size: int, heads: int, *, prefix: str) -> None: + super().__init__() + self.up = _head_sharded_linear( + hidden_size, heads, bias=True, prefix=f"{prefix}.up" + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, _ = self.up(x) + return torch.sigmoid(gate) + + +class VDNShortConv(nn.Module): + """Separable depthwise short conv (5x5 spatial per frame, then 5 taps + across frames) on the projections named in ``targets``; channels are + head-major, so TP takes a channel slice.""" + + def __init__(self, channels: int, targets: tuple[str, ...]) -> None: + super().__init__() + self.targets = tuple(targets) + k = SHORT_CONV_KERNEL + for name in self.targets: + setattr( + self, + f"{name}_sp", + nn.ParameterDict( + { + "weight": _make_param( + (channels, 1, k, k), dtype=_BF16, shard_dim=0 + ) + } + ), + ) + setattr( + self, + f"{name}_tm", + nn.ParameterDict( + {"weight": _make_param((channels, 1, k), dtype=_BF16, shard_dim=0)} + ), + ) + + def spatial( + self, + proj: str, + tokens: torch.Tensor, + num_frames: int, + frame_size: tuple[int, int], + heads: slice | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The 5x5 depthwise half on [F*S, H, d] tokens -> ([F, S, C], w_tm [C, 5]); + ``heads`` selects the weight channels of a head range.""" + n_heads, head_dim = tokens.shape[-2], tokens.shape[-1] + grid_h, grid_w = frame_size + channels = n_heads * head_dim + w_sp = getattr(self, f"{proj}_sp")["weight"] + w_tm = getattr(self, f"{proj}_tm")["weight"] + if heads is not None: + rows = slice(heads.start * head_dim, heads.stop * head_dim) + w_sp, w_tm = w_sp[rows], w_tm[rows] + # [F*S, H, d] read as channels_last [F, C, gh, gw]: cuDNN NHWC depthwise + volume = tokens.reshape(num_frames, grid_h, grid_w, channels).permute( + 0, 3, 1, 2 + ) + volume = F.conv2d(volume, w_sp, padding=SHORT_CONV_KERNEL // 2, groups=channels) + x = volume.permute(0, 2, 3, 1).reshape(num_frames, grid_h * grid_w, channels) + return x, w_tm.squeeze(1).to(x.dtype) + + +def _branch_norm(dim: int, eps: float = 1e-6) -> nn.RMSNorm: + # weight holder only; the arithmetic runs in the epilogue (fp32 second moment) + return nn.RMSNorm(dim, eps=eps, dtype=_BF16) + + +# -------------------------------------------------------------------------- +# The algorithm (eager, inference-only) +# -------------------------------------------------------------------------- + + +def _temporal_shift(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + # depthwise 5-tap conv over frames, zero padded, symmetric; x [F, S, C], w [C, 5] + k = SHORT_CONV_KERNEL + pad = k // 2 + xp = F.pad(x, (0, 0, 0, 0, pad, pad)) + out = None + for dt in range(k): + part = xp[dt : dt + x.shape[0]] * w[:, dt].view(1, 1, -1) + out = part if out is None else out + part + return out + + +def _activate(tokens: torch.Tensor, l2norm: bool) -> torch.Tensor: + x = F.silu(tokens) + return F.normalize(x, dim=-1, eps=1e-6).to(x.dtype) if l2norm else x + + +def linear_features( + tokens: torch.Tensor, + *, + proj: str, + conv: VDNShortConv | None, + num_frames: int | None, + frame_size: tuple[int, int] | None, + heads: slice | None = None, + frame_major: bool = False, + fused: bool = True, +) -> torch.Tensor: + """[N, H, d] raw projection -> [N, H, d] branch features: + [short conv ->] SiLU [-> L2 norm for q, k]. ``frame_major`` returns + [F, H, S, d] instead (the readout's bmm layout), written by the fused + kernels directly; the eager path permutes.""" + l2norm = proj != "v" + n_heads, head_dim = tokens.shape[-2], tokens.shape[-1] + if frame_major and (num_frames is None or frame_size is None): + raise ValueError("frame_major needs the (frames, height, width) grid") + if conv is not None and proj in conv.targets: + if frame_size is None or num_frames is None: + raise ValueError("the short conv needs the (frames, height, width) grid") + x, w_tm = conv.spatial(proj, tokens, num_frames, frame_size, heads=heads) + if fused and can_use_vdn_temporal_conv_act(x, n_heads, head_dim): + # one kernel: 5 taps + SiLU + L2 norm, the conv output never hits HBM + return vdn_temporal_conv_act( + x, w_tm, n_heads, head_dim, l2norm, frame_major=frame_major + ) + out = _activate(_temporal_shift(x, w_tm).reshape(-1, n_heads, head_dim), l2norm) + elif fused and can_use_vdn_silu_l2norm(tokens): + per_frame = frame_size[0] * frame_size[1] if frame_major else None + return vdn_silu_l2norm(tokens, l2norm, per_frame=per_frame) + else: + out = _activate(tokens, l2norm) + if frame_major: + per_frame = frame_size[0] * frame_size[1] + return out.view(num_frames, per_frame, n_heads, head_dim).permute(0, 2, 1, 3) + return out + + +def frame_statistics( + kf: torch.Tensor, + vf: torch.Tensor, + beta: torch.Tensor, + *, + a_fp32: bool, + prepared: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """kf, vf [F, H, S, d], beta [F, H, S] -> A [F, H, dk, dk] fp32 symmetric, + B [F, H, dv, dk] fp32; ``prepared`` carries the operands from + ``vdn_frame_stats_prep``. A is inverted downstream, so it needs fp32; + B enters the state linearly and takes bf16 tensor cores.""" + if prepared is not None: + kf, kf32, scaled32, vf_b = prepared + else: + kf = kf.contiguous() + vf_b = (vf * beta.unsqueeze(-1).to(vf.dtype)).contiguous() + kf32 = scaled32 = None + if a_fp32: + if kf32 is None: + kf32 = kf.float() + scaled32 = (kf32 * beta.unsqueeze(-1).float()).contiguous() + # TF32 keeps I + A well conditioned where bf16 does not; scoped to this matmul + prev = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + A = torch.matmul(scaled32.transpose(-1, -2), kf32) + finally: + torch.backends.cuda.matmul.allow_tf32 = prev + else: + A = torch.matmul( + (kf * beta.unsqueeze(-1).to(kf.dtype)).contiguous().transpose(-1, -2), kf + ).float() + A = 0.5 * (A + A.transpose(-1, -2)) + B = torch.matmul(vf_b.transpose(-1, -2), kf).float() + return A, B + + +def delta_factor_apply( + rule: str, + alpha: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + *, + tokens_per_frame: int, + fused: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """One frame's statistics -> (transition [F,H,dk,dk], injection [F,H,dv,dk]) fp32. + + vdn_solve (the released checkpoints): S' = (S diag(alpha) + B)(I + A)^-1, + exact Cholesky inverse. + sana_scaled: S' = (S diag(alpha))(I - c^2 A) + c B, c = 1/sqrt(S). + vdn_scaled: S' = (S diag(alpha) + c B)(I + c^2 A)^-1. + + ``fused`` takes the inverse and both products through one CUDA kernel + (``vdn_delta_factors``, fp32 / head_dim 128) with the same accuracy as the + Cholesky chain; anything the kernel does not cover falls back to eager. + """ + A32, B32 = A.float(), B.float() + eye = torch.eye(A32.shape[-1], device=A32.device, dtype=_FP32).expand_as(A32) + if rule == "sana_scaled": + inv_tokens = 1.0 / tokens_per_frame + transition = alpha.unsqueeze(-1) * (eye - inv_tokens * A32) + injection = math.sqrt(inv_tokens) * B32 + return transition, injection + if rule == "vdn_scaled": + inv_tokens = 1.0 / tokens_per_frame + A32 = A32 * inv_tokens + B32 = B32 * math.sqrt(inv_tokens) + elif rule != "vdn_solve": + raise ValueError(f"unknown delta rule {rule!r}") + if fused: + A32, B32, alpha32 = ( + A32.contiguous(), + B32.contiguous(), + alpha.float().contiguous(), + ) + if can_use_vdn_delta_factors(A32, B32, alpha32): + return vdn_delta_factors(A32, B32, alpha32) + chol = torch.linalg.cholesky(A32 + eye) + # (I+A)^-1 = L^-T L^-1: a batched trsm at 128x128 is far slower than the GEMM + linv = torch.linalg.solve_triangular(chol, eye, upper=False, left=True) + inv = linv.transpose(-1, -2) @ linv + transition = alpha.unsqueeze(-1) * inv + injection = B32 @ inv + return transition, injection + + +def run_scans( + transitions: torch.Tensor, + injections: torch.Tensor, + text_state: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """prefix[t] = frames 0..t, suffix[t] = frames t..F-1 (both fp32 + [F, H, dv, dk]); both start from ``text_state`` (or zero).""" + num_frames = transitions.shape[0] + start = ( + torch.zeros_like(injections[0]) + if text_state is None + else text_state.to(injections.dtype) + ) + prefix = torch.empty_like(injections) + suffix = torch.empty_like(injections) + state = start + for frame in range(num_frames): + torch.baddbmm(injections[frame], state, transitions[frame], out=prefix[frame]) + state = prefix[frame] + state = start + for frame in range(num_frames - 1, -1, -1): + torch.baddbmm(injections[frame], state, transitions[frame], out=suffix[frame]) + state = suffix[frame] + return prefix, suffix + + +def _compose_chunk( + transitions: torch.Tensor, injections: torch.Tensor, reverse: bool +) -> tuple[torch.Tensor, torch.Tensor]: + # fold each chunk's frames into one affine map S -> S @ M + C, batched over chunks + order = list(range(transitions.shape[0])) + if reverse: + order.reverse() + chunks, heads = transitions.shape[1], transitions.shape[2] + dk, dv = transitions.shape[-1], injections.shape[-2] + folded_t = transitions[order[0]] + folded_b = injections[order[0]] + for j in order[1:]: + step_t = transitions[j].view(chunks * heads, dk, dk) + folded_b = torch.baddbmm( + injections[j].view(chunks * heads, dv, dk), + folded_b.view(chunks * heads, dv, dk), + step_t, + ).view(chunks, heads, dv, dk) + folded_t = torch.bmm(folded_t.view(chunks * heads, dk, dk), step_t).view( + chunks, heads, dk, dk + ) + return folded_t, folded_b + + +@functools.lru_cache(maxsize=64) +def _boundary_frames( + num_frames: int, chunk: int, frame_offset: int, device: str +) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]: + # the gather reads prefix at chunk ends and suffix at chunk starts, on the offset grid + padded = frame_offset + num_frames + num_chunks = -(-padded // chunk) + ends = [ + min((c + 1) * chunk - 1, padded - 1) - frame_offset for c in range(num_chunks) + ] + starts = [c * chunk - frame_offset for c in range(num_chunks)] + dev = torch.device(device) + ends = [(f, c) for c, f in enumerate(ends) if f >= 0] + starts = [(f, c) for c, f in enumerate(starts) if f >= 0] + return ( + num_chunks, + torch.tensor([f for f, _ in ends], device=dev), + torch.tensor([c for _, c in ends], device=dev), + torch.tensor([f for f, _ in starts], device=dev), + torch.tensor([c for _, c in starts], device=dev), + ) + + +def run_boundary_scans( + transitions: torch.Tensor, + injections: torch.Tensor, + text_state: torch.Tensor | None, + *, + chunk: int, + frame_offset: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """``run_scans`` restricted to what the chunked gather reads: prefix at each + chunk's last frame, suffix at each chunk's first frame, zero elsewhere. + ``frame_offset`` is frame 0's position on the chunk grid (1 when the anchor + frames were dropped). Same fp32 math, the products re-associated.""" + if chunk <= 1: + return run_scans(transitions, injections, text_state) + num_frames, heads, dv, dk = injections.shape + num_chunks, ends, end_chunks, starts, start_chunks = _boundary_frames( + num_frames, chunk, frame_offset, str(injections.device) + ) + # identity / zero padding fills the leading offset and the partial last chunk + lead, tail = frame_offset, num_chunks * chunk - frame_offset - num_frames + eye = torch.eye(dk, device=transitions.device, dtype=transitions.dtype) + transitions = torch.cat( + [eye.expand(lead, heads, dk, dk), transitions, eye.expand(tail, heads, dk, dk)] + ) + injections = torch.cat( + [ + injections.new_zeros(lead, heads, dv, dk), + injections, + injections.new_zeros(tail, heads, dv, dk), + ] + ) + # frame-major so each composition step reads contiguous operands + by_frame_t = ( + transitions.view(num_chunks, chunk, heads, dk, dk).transpose(0, 1).contiguous() + ) + by_frame_b = ( + injections.view(num_chunks, chunk, heads, dv, dk).transpose(0, 1).contiguous() + ) + start = ( + torch.zeros(heads, dv, dk, dtype=injections.dtype, device=injections.device) + if text_state is None + else text_state.to(injections.dtype) + ) + # step c: the forward chain on chunk c and the reverse chain on chunk C-1-c + fwd_t, fwd_b = _compose_chunk(by_frame_t, by_frame_b, reverse=False) + rev_t, rev_b = _compose_chunk(by_frame_t, by_frame_b, reverse=True) + chunk_t = torch.stack([fwd_t, rev_t.flip(0)], dim=1) # [C, 2, H, dk, dk] + boundary = torch.stack([fwd_b, rev_b.flip(0)], dim=1) # [C, 2, H, dv, dk] + flat = boundary.view(num_chunks, 2 * heads, dv, dk) + state = torch.stack([start, start], dim=0).view(2 * heads, dv, dk) + for c in range(num_chunks): + flat[c].baddbmm_(state, chunk_t[c].view(2 * heads, dk, dk)) + state = flat[c] + prefix = torch.zeros( + num_frames, heads, dv, dk, dtype=injections.dtype, device=injections.device + ) + suffix = torch.zeros_like(prefix) + # step c holds chunk c's forward state and chunk C-1-c's reverse state + prefix.index_copy_(0, ends, boundary[end_chunks, 0]) + suffix.index_copy_(0, starts, boundary[num_chunks - 1 - start_chunks, 1]) + return prefix, suffix + + +@functools.lru_cache(maxsize=64) +def _gather_indices( + bounds: tuple[tuple[int, int], ...], num_frames: int, device: str +) -> tuple[torch.Tensor, ...]: + # cached: rebuilding from Python lists per block costs two synchronizing H2D copies + dev = torch.device(device) + last_before = torch.tensor([lo for lo, _ in bounds], device=dev) - 1 + first_after = torch.tensor([hi for _, hi in bounds], device=dev) + 1 + return ( + last_before, + first_after, + last_before.clamp(min=0), + first_after.clamp(max=num_frames - 1), + last_before >= 0, + first_after < num_frames, + torch.arange(num_frames, device=dev), + ) + + +def gather_linear_state( + prefix: torch.Tensor, + suffix: torch.Tensor, + alpha: torch.Tensor, + bounds: list[tuple[int, int]], + *, + bridge: str, + text_state: torch.Tensor | None, + out_dtype: torch.dtype, + fused: bool = True, +) -> torch.Tensor: + """Everything OUTSIDE the softmax window of frame t, decayed to t: + prefix[lo-1] * prod_{u=lo..t} alpha_u + suffix[hi+1] * prod_{u=t..hi} alpha_u. + Out-of-range sides read the text state (the scans' virtual start) when one + was given, else contribute nothing. -> [F, H, dv, dk] in ``out_dtype``.""" + num_frames = prefix.shape[0] + ( + last_before, + first_after, + before_idx, + after_idx, + has_before, + has_after, + frames, + ) = _gather_indices(tuple(bounds), num_frames, str(prefix.device)) + if fused and can_use_vdn_gather_linear_state(prefix): + return vdn_gather_linear_state( + prefix, + suffix, + alpha, + text_state, + before_idx=before_idx, + after_idx=after_idx, + has_before=has_before, + has_after=has_after, + bridge_before=(last_before + 1).clamp(min=0), + bridge_after=first_after.clamp(max=num_frames), + bridge=bridge == "alpha", + out_dtype=out_dtype, + ) + + state_before = prefix[before_idx] + state_after = suffix[after_idx] + if text_state is not None: + ts = text_state.to(state_before.dtype) + state_before = torch.where(has_before.view(-1, 1, 1, 1), state_before, ts) + state_after = torch.where(has_after.view(-1, 1, 1, 1), state_after, ts) + if bridge == "alpha": + log_alpha = torch.log(alpha.clamp_min(1e-12)) + log_prefix = torch.cat([torch.zeros_like(log_alpha[:1]), log_alpha.cumsum(0)]) + # an out-of-range side decays the text state over [0..t] or [t..F-1] + bridge_before = (last_before + 1).clamp(min=0) + bridge_after = first_after.clamp(max=num_frames) + alpha_from_before = torch.exp( + log_prefix[frames + 1] - log_prefix[bridge_before] + ) + alpha_from_after = torch.exp(log_prefix[bridge_after] - log_prefix[frames]) + # alpha is per KEY channel: broadcast over dv, not dk + state_before = state_before * alpha_from_before.unsqueeze(2) + state_after = state_after * alpha_from_after.unsqueeze(2) + elif bridge != "none": + raise ValueError(f"unknown bridge {bridge!r}") + if text_state is not None: + out = state_before + state_after + else: + out = state_before * has_before.view( + -1, 1, 1, 1 + ) + state_after * has_after.view(-1, 1, 1, 1) + return out.to(out_dtype) + + +def linear_epilogue( + readout: torch.Tensor, norm_weight: torch.Tensor, gate: torch.Tensor, eps: float +) -> torch.Tensor: + """readout [F, H, S, dv] -> RMSNorm over dv -> * gate [F*S, H, dv] -> [F*S, H*dv].""" + ms = ( + torch.linalg.vector_norm(readout, dim=-1, keepdim=True, dtype=_FP32).pow(2) + / (readout.shape[-1]) + ) + normed = ( + readout + * torch.rsqrt(ms + eps).to(readout.dtype) + * norm_weight.to(readout.dtype) + ) + frames, heads, per_frame, dim = normed.shape + rows = frames * per_frame + return normed.permute(0, 2, 1, 3).reshape(rows, heads * dim) * gate.reshape( + rows, heads * dim + ) + + +# -------------------------------------------------------------------------- +# The module +# -------------------------------------------------------------------------- + + +class MiniMaxH3VDNLinearBranch(nn.Module): + """VDN's BidirectionalLinearBranch on SGLang's TP-local head shard.""" + + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + hybrid: VDNHybridAttentionArchConfig, + *, + local_heads: int, + prefix: str = "linear_attention", + ) -> None: + super().__init__() + if hybrid.linear_head_dim != arch.attention_head_dim: + # the branch reads the attention projections, so the head dims must agree + raise ValueError( + f"hybrid_attention.linear_head_dim={hybrid.linear_head_dim} != " + f"attention_head_dim={arch.attention_head_dim}" + ) + self.hybrid = hybrid + self.local_heads = local_heads + self.head_dim = arch.attention_head_dim + # tests flip this to compare the fused Triton stages with the eager chain + self.fused_kernels = True + hidden = arch.hidden_size + channels = local_heads * self.head_dim + self.short_conv = ( + VDNShortConv(channels, hybrid.short_conv) if hybrid.short_conv else None + ) + heads = arch.num_attention_heads + self.alpha = VDNFrameAlpha( + hidden, heads, local_heads, self.head_dim, prefix=f"{prefix}.alpha" + ) + self.beta_proj = _head_sharded_linear( + hidden, heads, bias=False, prefix=f"{prefix}.beta_proj" + ) + self.output_gate = VDNOutputGate( + hidden, heads, local_heads, self.head_dim, prefix=f"{prefix}.output_gate" + ) + self.norm = _branch_norm(self.head_dim) + + # ---- pieces the attention module computes on the row shard (Ulysses) ---- + + def beta(self, x: torch.Tensor) -> torch.Tensor: + """x [T, hidden] -> beta [T, H_local] (sigmoid).""" + beta, _ = self.beta_proj(x) + return torch.sigmoid(beta) + + # ---- the text state ----------------------------------------------------- + + def text_statistics( + self, + text_k_raw: torch.Tensor, + text_v_raw: torch.Tensor, + text_beta: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + """(A [1, H, dk, dk], B [1, H, dv, dk]) fp32 of the prompt rows, no conv.""" + length = text_k_raw.shape[0] + heads, head_dim = text_k_raw.shape[1], self.head_dim + key = linear_features( + text_k_raw, + proj="k", + conv=None, + num_frames=None, + frame_size=None, + fused=self.fused_kernels, + ) + value = linear_features( + text_v_raw, + proj="v", + conv=None, + num_frames=None, + frame_size=None, + fused=self.fused_kernels, + ) + key = key.view(1, length, heads, head_dim).permute(0, 2, 1, 3) + value = value.view(1, length, heads, head_dim).permute(0, 2, 1, 3) + beta = text_beta.view(1, length, heads).permute(0, 2, 1) + A, B = frame_statistics(key, value, beta, a_fp32=self.hybrid.a_fp32) + return A, B, length + + def text_state( + self, + text_k_raw: torch.Tensor, + text_v_raw: torch.Tensor, + text_beta: torch.Tensor, + ) -> torch.Tensor: + """S_text [H, dv, dk] fp32: the prompt written into a zero state as one + delta-rule chunk, scaled by TEXT_STATE_SCALE.""" + A, B, length = self.text_statistics(text_k_raw, text_v_raw, text_beta) + heads, head_dim = A.shape[1], self.head_dim + ones = torch.ones(1, heads, head_dim, device=A.device, dtype=_FP32) + _, injection = delta_factor_apply( + self.hybrid.delta_rule, + ones, + A, + B, + tokens_per_frame=length, + fused=self.fused_kernels, + ) + return TEXT_STATE_SCALE * injection[0] + + # ---- the branch -------------------------------------------------------- + + def forward( + self, + *, + q_raw: torch.Tensor, + k_raw: torch.Tensor, + v_raw: torch.Tensor, + beta: torch.Tensor, + gate: torch.Tensor, + frame_mean: torch.Tensor, + layout: VDNH3Layout, + text_k_raw: torch.Tensor | None = None, + text_v_raw: torch.Tensor | None = None, + text_beta: torch.Tensor | None = None, + heads: slice | None = None, + ) -> torch.Tensor: + """Linear readout of the video rows, [V, H * d] in q's dtype. + + q/k/v/beta/gate are the video rows' raw (pre-norm, pre-RoPE) values on + this rank's heads; ``heads`` is the head range of the full sequence a + Ulysses rank processes, applied to the per-head parameters. Under + anchor_frames == "both" frames 0 and F-1 read zero. + """ + hybrid = self.hybrid + num_frames, per_frame = layout.num_frames, layout.tokens_per_frame + bounds = hybrid.window_bounds(num_frames) + text_state = None + text_stats = None + if hybrid.enable_text_state: + if text_k_raw is None or text_v_raw is None or text_beta is None: + raise ValueError("enable_text_state needs the prompt rows' k/v/beta") + if text_k_raw.shape[0] > 0: + if hybrid.delta_rule == "vdn_solve": + # vdn_solve ignores tokens_per_frame, so the prompt joins the batch + A_text, B_text, _ = self.text_statistics( + text_k_raw, text_v_raw, text_beta + ) + text_stats = (A_text, B_text) + else: + text_state = self.text_state(text_k_raw, text_v_raw, text_beta) + + skip_ends = hybrid.anchor_frames == "both" + n_heads = q_raw.shape[1] + if not skip_ends: + return self._readout( + q_raw, + k_raw, + v_raw, + beta, + gate, + frame_mean, + num_frames, + per_frame, + bounds, + layout.frame_size, + text_state, + heads, + text_stats=text_stats, + ) + out = q_raw.new_empty(num_frames * per_frame, n_heads * self.head_dim) + if num_frames <= 2: + return out.zero_() + inner = slice(per_frame, (num_frames - 1) * per_frame) + readout = self._readout( + q_raw[inner], + k_raw[inner], + v_raw[inner], + beta[inner], + gate[inner], + frame_mean[1:-1], + num_frames - 2, + per_frame, + [(lo - 1, hi - 1) for lo, hi in bounds[1 : num_frames - 1]], + layout.frame_size, + text_state, + heads, + frame_offset=1, + text_stats=text_stats, + ) + out[:per_frame].zero_() + out[(num_frames - 1) * per_frame :].zero_() + out[inner] = readout + return out + + def _readout( + self, + q_raw: torch.Tensor, + k_raw: torch.Tensor, + v_raw: torch.Tensor, + beta: torch.Tensor, + gate: torch.Tensor, + frame_mean: torch.Tensor, + num_frames: int, + per_frame: int, + bounds: list[tuple[int, int]], + frame_size: tuple[int, int], + text_state: torch.Tensor | None, + heads: slice | None, + frame_offset: int = 0, + text_stats: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + n_heads, head_dim = q_raw.shape[1], self.head_dim + shape = (num_frames, per_frame, n_heads, head_dim) + fused = self.fused_kernels + features = functools.partial( + linear_features, + conv=self.short_conv, + num_frames=num_frames, + frame_size=frame_size, + heads=heads, + fused=fused, + ) + query_by_frame = features(q_raw, proj="q", frame_major=True) + key = features(k_raw, proj="k") + value = features(v_raw, proj="v") + key_by_frame = key.view(shape).permute(0, 2, 1, 3) + value_by_frame = value.view(shape).permute(0, 2, 1, 3) + beta_by_frame = beta.view(num_frames, per_frame, n_heads).permute(0, 2, 1) + prepared = ( + vdn_frame_stats_prep(key, value, beta, num_frames, per_frame) + if fused and self.hybrid.a_fp32 and can_use_vdn_frame_stats_prep(key, value) + else None + ) + A, B = frame_statistics( + key_by_frame, + value_by_frame, + beta_by_frame, + a_fp32=self.hybrid.a_fp32, + prepared=prepared, + ) + del prepared + alpha = self.alpha(frame_mean, heads=heads) + if text_stats is not None: + # the prompt leads as a virtual frame; alpha 1 since its old state is zero + A = torch.cat([text_stats[0], A]) + B = torch.cat([text_stats[1], B]) + alpha_all = torch.cat([alpha.new_ones((1,) + alpha.shape[1:]), alpha]) + else: + alpha_all = alpha + transitions, injections = delta_factor_apply( + self.hybrid.delta_rule, + alpha_all, + A, + B, + tokens_per_frame=per_frame, + fused=self.fused_kernels, + ) + if text_stats is not None: + text_state = TEXT_STATE_SCALE * injections[0] + transitions, injections = transitions[1:], injections[1:] + prefix, suffix = run_boundary_scans( + transitions, + injections, + text_state, + chunk=self.hybrid.chunk, + frame_offset=frame_offset, + ) + del transitions, injections + linear_state = gather_linear_state( + prefix, + suffix, + alpha, + bounds, + bridge=self.hybrid.bridge, + text_state=text_state, + out_dtype=q_raw.dtype, + fused=fused, + ) + del prefix, suffix + readout = torch.matmul(query_by_frame, linear_state.transpose(-1, -2)) + if fused and can_use_vdn_linear_epilogue(readout): + return vdn_linear_epilogue(readout, self.norm.weight, gate, self.norm.eps) + return linear_epilogue(readout, self.norm.weight, gate, self.norm.eps) + + +__all__ = [ + "MiniMaxH3VDNLinearBranch", + "TEXT_STATE_SCALE", + "VDNFrameAlpha", + "VDNH3Layout", + "VDNOutputGate", + "VDNShortConv", + "VDNSoftmaxGate", + "delta_factor_apply", + "frame_statistics", + "gather_linear_state", + "linear_epilogue", + "linear_features", + "run_boundary_scans", + "run_scans", + "vdn_h3_layout_from_packed", +] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn_attention.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn_attention.py new file mode 100644 index 000000000..3413084ff --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn_attention.py @@ -0,0 +1,647 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 hybrid attention inside a MiniMax-H3 DiT block: per-head softmax gate, +Video Delta linear branch and its output projection, the Ulysses exchange that +shards both branches by head, and the request-static attention metadata. +``MiniMaxH3Attention`` owns one instance as ``hybrid`` and hands it raw q/k/v.""" + +from __future__ import annotations + +import functools +import logging +from typing import TYPE_CHECKING, Any, Callable, Mapping + +import torch +from torch import nn + +from sglang.kernels.ops.diffusion import ( + fused_qknorm_rope_out_of_place, + usp_merge_heads, +) +from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTArchConfig +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_ctx, + get_ulysses_ctx, +) +from sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3 import ( + HybridWindowAttentionH3Metadata, + HybridWindowAttentionH3MetadataBuilder, +) +from sglang.multimodal_gen.runtime.layers.linear import RowParallelLinear +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, +) +from sglang.multimodal_gen.runtime.layers.usp import _a2a_staging_buffer +from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import ( + MiniMaxH3VDNLinearBranch, + VDNSoftmaxGate, + vdn_h3_layout_from_packed, +) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + eager_on_graph, +) + +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3Attention + +logger = logging.getLogger(__name__) +_FP32_DTYPE = torch.float32 +_BF16_DTYPE = torch.bfloat16 + +# one side stream per process and device: under Ulysses the linear readout and +# its a2a run on it while FA4 holds the current stream +_LINEAR_STREAMS: dict[int, torch.cuda.Stream] = {} + + +def _linear_branch_stream(device: torch.device) -> torch.cuda.Stream: + index = device.index if device.index is not None else torch.cuda.current_device() + stream = _LINEAR_STREAMS.get(index) + if stream is None: + # high priority: the readout's small kernels fill the SM slots FA4's + # persistent CTAs leave at their tails + stream = torch.cuda.Stream(device=device, priority=-1) + _LINEAR_STREAMS[index] = stream + return stream + + +class MiniMaxH3VDNHybridAttention(nn.Module): + """out = to_out(gate_sm * window_softmax(q, k, v)) + to_out_linear(branch(q, k, v)).""" + + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + *, + prefix: str, + local_heads: int, + ) -> None: + super().__init__() + hybrid = arch.hybrid_attention + self.softmax_gate: VDNSoftmaxGate | None = None + if hybrid.enable_softmax_gate: + self.softmax_gate = VDNSoftmaxGate( + arch.hidden_size, + arch.num_attention_heads, + prefix=f"{prefix}.softmax_gate", + ) + self.linear_attention = MiniMaxH3VDNLinearBranch( + arch, hybrid, local_heads=local_heads, prefix=f"{prefix}.linear_attention" + ) + self.to_out_linear = RowParallelLinear( + arch.num_attention_heads * arch.attention_head_dim, + arch.hidden_size, + bias=False, + input_is_parallel=True, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + prefix=f"{prefix}.to_out_linear", + ) + + @classmethod + def build( + cls, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + *, + prefix: str, + local_heads: int, + ) -> MiniMaxH3VDNHybridAttention | None: + """None for the dense model and the token refiner (VDN converts the DiT blocks only).""" + if arch.hybrid_attention is None or not prefix.startswith("blocks."): + return None + return cls( + arch, quant_config, prefix=f"{prefix}.hybrid", local_heads=local_heads + ) + + def forward( + self, + attention: MiniMaxH3Attention, + x: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + rope_cache: tuple[torch.Tensor, torch.Tensor] | None, + cu_seqlens: torch.Tensor, + cu_seqlens_host: tuple[int, ...] | None, + max_seqlen: int, + ulysses_active: bool, + ring_active: bool, + ) -> torch.Tensor: + """out = to_out(gate_sm * window_softmax) + to_out_linear(branch); gates and + beta are row-local and computed before the core exchanges rows for heads.""" + if ring_active: + raise NotImplementedError( + "VDN-H3 hybrid attention does not support ring parallelism" + ) + total = x.shape[0] + softmax_gate = self.softmax_gate(x) if self.softmax_gate is not None else None + beta = self.linear_attention.beta(x) + gate_hidden, _ = self.linear_attention.output_gate.down(x) + attention_core = ( + _hybrid_attention_core_bcg + if attention.bcg_breakpoint + else _minimax_h3_hybrid_attention_core_impl + ) + softmax_out, linear_out = attention_core( + attention, + x, + q, + k, + v, + softmax_gate, + beta, + gate_hidden, + rope_cache=rope_cache, + cu_seqlens=cu_seqlens, + cu_seqlens_host=cu_seqlens_host, + max_seqlen=max_seqlen, + ulysses_active=ulysses_active, + ) + out, _ = attention.out_proj(softmax_out.reshape(total, -1)) + if linear_out is not None: + linear_proj, _ = self.to_out_linear(linear_out) + if linear_proj.shape[0] == total: + out.add_(linear_proj) + else: + # single-rank path: the readout covers the video rows only + layout = get_forward_context().attn_metadata.layout + out[layout.video_start : layout.video_end].add_(linear_proj) + return out + + +def _vdn_frame_partial_sums( + x: torch.Tensor, + *, + row_start: int, + video_start: int, + video_end: int, + num_frames: int, + tokens_per_frame: int, +) -> torch.Tensor: + # fp32 [F, hidden] sums of this rank's video rows; whole frames as one reduction + hidden = x.shape[-1] + sums = torch.zeros(num_frames, hidden, dtype=_FP32_DTYPE, device=x.device) + lo = max(row_start, video_start) + hi = min(row_start + x.shape[0], video_end) + if lo >= hi: + return sums + rows = x[lo - row_start : hi - row_start] + first_frame, offset = divmod(lo - video_start, tokens_per_frame) + lead = (tokens_per_frame - offset) % tokens_per_frame + lead = min(lead, rows.shape[0]) + if lead: + sums[first_frame] += rows[:lead].sum(0, dtype=_FP32_DTYPE) + first_frame += 1 + full = (rows.shape[0] - lead) // tokens_per_frame + if full: + sums[first_frame : first_frame + full] = ( + rows[lead : lead + full * tokens_per_frame] + .view(full, tokens_per_frame, hidden) + .sum(1, dtype=_FP32_DTYPE) + ) + tail = lead + full * tokens_per_frame + if tail < rows.shape[0]: + sums[first_frame + full] += rows[tail:].sum(0, dtype=_FP32_DTYPE) + return sums + + +def _vdn_a2a_rows_to_heads( + field: torch.Tensor, + *, + ulysses_ws: int, + role: str, + process_group: torch.distributed.ProcessGroup, +) -> tuple[torch.distributed.Work, torch.Tensor]: + # [L, H, d] row shard -> contiguous [S, H / ws, d] of this rank's heads + rows, total_heads, head_dim = field.shape + local_heads = total_heads // ulysses_ws + send = _a2a_staging_buffer( + role + "_send", + (ulysses_ws, rows, local_heads, head_dim), + field.dtype, + field.device, + ) + send.copy_(field.view(rows, ulysses_ws, local_heads, head_dim).permute(1, 0, 2, 3)) + recv = _a2a_staging_buffer( + role + "_recv", + (ulysses_ws * rows, local_heads, head_dim), + field.dtype, + field.device, + ) + work = torch.distributed.all_to_all_single( + recv, send, group=process_group, async_op=True + ) + return work, recv + + +def _vdn_a2a_heads_to_rows( + out: torch.Tensor, + *, + ulysses_ws: int, + role: str, + process_group: torch.distributed.ProcessGroup, +) -> tuple[torch.distributed.Work, torch.Tensor]: + # [S, H / ws, d] -> [ws, L, H / ws, d] source-rank major; _vdn_merge_heads after wait + seq_len, local_heads, head_dim = out.shape + rows = seq_len // ulysses_ws + recv = _a2a_staging_buffer( + role + "_recv", (ulysses_ws, rows, local_heads, head_dim), out.dtype, out.device + ) + work = torch.distributed.all_to_all_single( + recv, out.contiguous(), group=process_group, async_op=True + ) + return work, recv + + +def _vdn_merge_heads(recv: torch.Tensor) -> torch.Tensor: + # [ws, L, h, d] -> [L, ws * h, d]; rank-major heads are the global head order + ulysses_ws, rows, local_heads, head_dim = recv.shape + merged = usp_merge_heads(recv.view(ulysses_ws, rows, 1, local_heads, head_dim)) + return merged.reshape(rows, ulysses_ws * local_heads, head_dim) + + +def _vdn_window_softmax( + attention: MiniMaxH3Attention, + meta: HybridWindowAttentionH3Metadata, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + softmax_gate: torch.Tensor | None, + rope_cache: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.Tensor, + cu_seqlens_host: tuple[int, ...] | None, + max_seqlen: int, +) -> torch.Tensor: + # the branch keeps reading the raw q/k, so norm + RoPE write copies + cos_sin_cache, positions = rope_cache + if attention._use_fused_qknorm_rope and not torch.compiler.is_compiling(): + q_sm = torch.empty(q.shape, dtype=q.dtype, device=q.device) + k_sm = torch.empty(k.shape, dtype=k.dtype, device=k.device) + fused_qknorm_rope_out_of_place( + q, + k, + q_sm, + k_sm, + attention.q_norm.weight, + attention.k_norm.weight, + cos_sin_cache, + positions, + is_neox=True, + eps=attention.q_norm.eps, + head_dim=attention.head_dim, + rope_dim=cos_sin_cache.shape[-1], + round_norm_before_rope=True, + ) + else: + from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( + _apply_qk_norm, + _apply_rope_qk, + ) + + q_sm, k_sm = _apply_qk_norm( + q.clone(), k.clone(), attention.q_norm, attention.k_norm, attention.head_dim + ) + q_sm, k_sm = _apply_rope_qk(q_sm, k_sm, cos_sin_cache, positions) + return attention._attention_impl.forward_varlen( + q_sm, + k_sm, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + attn_metadata=meta, + softmax_gate=softmax_gate, + ) + + +def _vdn_linear_readout( + attention: MiniMaxH3Attention, + meta: HybridWindowAttentionH3Metadata, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + beta: torch.Tensor, + linear_gate: torch.Tensor, + frame_mean: torch.Tensor, + head_range: slice | None, +) -> torch.Tensor: + # video rows in, [V, h * d] out; text rows seed the state + layout = meta.layout + video = slice(layout.video_start, layout.video_end) + text = slice(0, layout.text_len) + return attention.hybrid.linear_attention( + q_raw=q[video], + k_raw=k[video], + v_raw=v[video], + beta=beta[video], + gate=linear_gate[video], + frame_mean=frame_mean, + layout=layout, + text_k_raw=k[text], + text_v_raw=v[text], + text_beta=beta[text], + heads=head_range, + ) + + +def _minimax_h3_hybrid_attention_core_impl( + attention: MiniMaxH3Attention, + x: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_gate: torch.Tensor | None, + beta: torch.Tensor, + gate_hidden: torch.Tensor, + *, + rope_cache: tuple[torch.Tensor, torch.Tensor] | None, + cu_seqlens: torch.Tensor, + cu_seqlens_host: tuple[int, ...] | None, + max_seqlen: int, + ulysses_active: bool, +) -> tuple[torch.Tensor, torch.Tensor | None]: + meta = get_forward_context().attn_metadata + if not isinstance(meta, HybridWindowAttentionH3Metadata): + raise RuntimeError( + "VDN-H3 hybrid attention needs HybridWindowAttentionH3Metadata in the " + "forward context; the MiniMax-H3 denoising stage installs it per request " + f"(got {type(meta).__name__})." + ) + layout = meta.layout + softmax = functools.partial( + _vdn_window_softmax, + attention, + meta, + cu_seqlens=cu_seqlens, + cu_seqlens_host=cu_seqlens_host, + max_seqlen=max_seqlen, + ) + if not ulysses_active: + if rope_cache is None: + raise RuntimeError("VDN-H3 hybrid attention requires the RoPE cache") + softmax_out = softmax(q, k, v, softmax_gate=softmax_gate, rope_cache=rope_cache) + if meta.full_cover: + return softmax_out, None + frame_mean = ( + x[layout.video_start : layout.video_end] + .view(layout.num_frames, layout.tokens_per_frame, x.shape[-1]) + .mean(dim=1, dtype=_FP32_DTYPE) + ) + readout = _vdn_linear_readout( + attention, + meta, + q, + k, + v, + beta=beta, + linear_gate=attention.hybrid.linear_attention.output_gate.up_gate( + gate_hidden + ), + frame_mean=frame_mean, + head_range=None, + ) + return softmax_out, readout + + return _vdn_ulysses_hybrid_core( + attention, + meta, + x, + q, + k, + v, + softmax_gate=softmax_gate, + beta=beta, + gate_hidden=gate_hidden, + softmax=softmax, + ) + + +def _vdn_ulysses_hybrid_core( + attention: MiniMaxH3Attention, + meta: HybridWindowAttentionH3Metadata, + x: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + softmax_gate: torch.Tensor | None, + beta: torch.Tensor, + gate_hidden: torch.Tensor, + softmax: Callable[..., torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor | None]: + from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group + + layout = meta.layout + if meta.rope_cache_full is None: + raise RuntimeError("VDN-H3 under Ulysses needs the full-sequence RoPE cache") + sp_group = get_sp_group() + process_group = sp_group.ulysses_group + ulysses_ws, ulysses_rank = get_ulysses_ctx() + local_rows, head_dim = x.shape[0], q.shape[2] + local_heads = q.shape[1] // ulysses_ws + seq_len = local_rows * ulysses_ws + + # the q/k/v exchange is in flight while the frame sums and the gate hidden go out + inflight = [ + _vdn_a2a_rows_to_heads( + field, + ulysses_ws=ulysses_ws, + role=f"vdn_{name}", + process_group=process_group, + ) + for name, field in (("q", q), ("k", k), ("v", v)) + ] + frame_sums = _vdn_frame_partial_sums( + x, + row_start=ulysses_rank * local_rows, + video_start=layout.video_start, + video_end=layout.video_end, + num_frames=layout.num_frames, + tokens_per_frame=layout.tokens_per_frame, + ) + frame_work = torch.distributed.all_reduce( + frame_sums, group=sp_group.device_group, async_op=True + ) + # the per-head scalars (beta, softmax gate) ride one more async field + scalars = [beta] if softmax_gate is None else [beta, softmax_gate] + inflight.append( + _vdn_a2a_rows_to_heads( + torch.stack(scalars, dim=-1), + ulysses_ws=ulysses_ws, + role="vdn_scalars", + process_group=process_group, + ) + ) + head_range = slice(ulysses_rank * local_heads, (ulysses_rank + 1) * local_heads) + linear_gate = attention.hybrid.linear_attention.output_gate.up_gate( + sp_group.all_gather(gate_hidden.contiguous(), dim=0), heads=head_range + ) + + for work, _ in inflight: + work.wait() + q, k, v, scalars = (recv for _, recv in inflight) + beta = scalars[..., 0] + if softmax_gate is not None: + softmax_gate = scalars[..., 1] + softmax = functools.partial( + softmax, q, k, v, softmax_gate=softmax_gate, rope_cache=meta.rope_cache_full + ) + if meta.full_cover: + softmax_out = softmax() + frame_work.wait() + return _vdn_return_to_rows( + softmax_out, None, ulysses_ws=ulysses_ws, process_group=process_group + ) + + def linear_branch() -> torch.Tensor: + frame_work.wait() + readout = _vdn_linear_readout( + attention, + meta, + q, + k, + v, + beta=beta, + linear_gate=linear_gate, + frame_mean=frame_sums / layout.tokens_per_frame, + head_range=head_range, + ) + # rows go back to their owners: pad the non-video rows with zeros + linear_out = q.new_zeros(seq_len, local_heads, head_dim) + linear_out[layout.video_start : layout.video_end] = readout.view( + -1, local_heads, head_dim + ) + return linear_out + + a2a_back = functools.partial( + _vdn_a2a_heads_to_rows, ulysses_ws=ulysses_ws, process_group=process_group + ) + # the linear readout and its a2a run on the side stream while FA4 (issued + # first) holds the current one; both read the exchanged q/k/v, frame sums and + # gate, so the side stream waits for them and the current stream joins before the merge + main_stream = torch.cuda.current_stream(q.device) + side_stream = _linear_branch_stream(q.device) + side_stream.wait_stream(main_stream) + softmax_out = softmax() + softmax_work, softmax_recv = a2a_back(softmax_out, role="vdn_out0") + with torch.cuda.stream(side_stream): + linear_work, linear_recv = a2a_back(linear_branch(), role="vdn_out1") + main_stream.wait_stream(side_stream) + softmax_work.wait() + linear_work.wait() + merged_softmax = _vdn_merge_heads(softmax_recv) + merged_linear = _vdn_merge_heads(linear_recv) + return merged_softmax, merged_linear.reshape(merged_linear.shape[0], -1) + + +def _vdn_return_to_rows( + softmax_out: torch.Tensor, + linear_out: torch.Tensor | None, + *, + ulysses_ws: int, + process_group: torch.distributed.ProcessGroup, +) -> tuple[torch.Tensor, torch.Tensor | None]: + # [S, h, d] per branch -> ([L, H, d], [L, H * d] or None), both trips in flight together + branch_outputs = [out for out in (softmax_out, linear_out) if out is not None] + inflight = [ + _vdn_a2a_heads_to_rows( + out, ulysses_ws=ulysses_ws, role=f"vdn_out{i}", process_group=process_group + ) + for i, out in enumerate(branch_outputs) + ] + merged = [] + for work, recv in inflight: + work.wait() + merged.append(_vdn_merge_heads(recv)) + linear_rows = ( + merged[1].reshape(merged[1].shape[0], -1) if linear_out is not None else None + ) + return merged[0], linear_rows + + +_hybrid_attention_core_bcg = eager_on_graph(True)( + _minimax_h3_hybrid_attention_core_impl +) + + +def prepare_hybrid_attention_metadata( + *, + model, + packed: Mapping[str, torch.Tensor], + latent_shape: tuple[int, int, int], + server_args, + device: torch.device, +) -> Callable[[int], Any] | None: + """Request-static metadata (window plan, packed layout, full-sequence RoPE + cache under Ulysses) for every step and block; None for other backends.""" + model._resolve_attention_backend_once() + if ( + model._resolved_attention_backend + is not AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3 + ): + return None + hybrid = model.arch.hybrid_attention + if hybrid is None: + raise ValueError( + "--attention-backend hybrid_window_attn_h3 needs a VDN-H3 checkpoint " + "(transformer/config.json with hybrid_attention); this checkpoint has " + "no linear branch. Use --attention-backend fa for MiniMax-H3." + ) + from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( + _rope_cos_sin_cache, + ) + + config = server_args.attention_backend_config or {} + max_gather_rows = int(config.get("vdn_max_gather_rows", 200_000)) + latent_t, latent_h, latent_w = latent_shape + layout = vdn_h3_layout_from_packed( + packed, latent_t=latent_t, latent_h=latent_h, latent_w=latent_w + ) + rope_cache_full = None + ulysses_ws, _ = get_ulysses_ctx() + ring_ws, _ = get_ring_ctx() + if ring_ws > 1: + raise ValueError("VDN-H3 does not support ring parallelism") + if ulysses_ws > 1: + # QK-norm + RoPE run on the head shard after the all-to-all: full-sequence cache + with torch.inference_mode(): + img_position_ids = ( + packed["img_position_ids"][None].to(torch.float32).to(device) + ) + rope_freqs = model.rope(img_position_ids) + rope_cache_full = ( + _rope_cos_sin_cache(rope_freqs, dtype=torch.bfloat16), + torch.arange(layout.seq_len, device=device, dtype=torch.long), + ) + metadata = HybridWindowAttentionH3MetadataBuilder().build( + layout=layout, + hybrid=hybrid, + device=device, + rope_cache_full=rope_cache_full, + max_gather_rows=max_gather_rows, + ) + logger.info( + "VDN-H3 hybrid attention: frames=%d tokens/frame=%d text=%d " + "used=%d/%d chunk=%d radius=%d anchors=%s full_cover=%s", + layout.num_frames, + layout.tokens_per_frame, + layout.text_len, + layout.used, + layout.seq_len, + hybrid.chunk, + hybrid.radius, + hybrid.anchor_frames, + metadata.full_cover, + ) + + def build(step_index: int): + return metadata + + return build + + +__all__ = ["MiniMaxH3VDNHybridAttention", "prepare_hybrid_attention_metadata"] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_vdn_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_vdn_pipeline.py new file mode 100644 index 000000000..0587a1df8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_vdn_pipeline.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +from sglang.multimodal_gen.runtime.pipelines.minimax_h3_pipeline import ( + MiniMaxH3Pipeline, +) + + +class VDNH3Pipeline(MiniMaxH3Pipeline): + """VDN-H3 on the MiniMax-H3 pipeline: the model overlay materializes a + base-H3 layout (LoRAs prefused, linear branch attached), so only the DiT + blocks' attention and its backend differ.""" + + pipeline_name = "VDNH3Pipeline" + default_model_subfolder = None + + +EntryClass = [VDNH3Pipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py index 5d5901b24..ab65f65c5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py @@ -749,6 +749,18 @@ class MiniMaxH3DenoisingStage(DenoisingStage): server_args=server_args, device=device, ) + if build_vsa_h3_step_metadata is None: + from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import ( + prepare_hybrid_attention_metadata, + ) + + build_vsa_h3_step_metadata = prepare_hybrid_attention_metadata( + model=model, + packed=packed, + latent_shape=(ctx.latent_t, ctx.latent_h, ctx.latent_w), + server_args=server_args, + device=device, + ) positive = MiniMaxH3DenoiseBranch( packed=packed, text_embeddings=emb["hidden_states"], diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index b3b5f52d4..f1878f8f0 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -290,6 +290,52 @@ class _VideoSparseAttentionH3BackendResolver(_CudaAttentionBackendResolver): ) from e +class _HybridWindowAttentionH3BackendResolver(_CudaAttentionBackendResolver): + backend = AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3 + + # the window rides FlashAttention varlen: FA4 on SM100 / SM103 / SM120, FA3 on + # SM90; SM80 / SM86 / SM89 run FA3's Sm80 mainloop (FA2-class throughput) + supported_capabilities = { + (8, 0), + (8, 6), + (8, 9), + (9, 0), + (10, 0), + (10, 3), + (12, 0), + } + + @classmethod + def resolve(cls, platform) -> str: + capability = platform.get_device_capability() + capability_tuple = ( + (capability.major, capability.minor) if capability is not None else None + ) + if capability_tuple not in cls.supported_capabilities: + found = capability.as_version_str() if capability else "unknown" + raise ValueError( + "hybrid_window_attn_h3 (VDN-H3) needs compute capability 8.0 / " + "8.6 / 8.9 (Ampere, Ada), 9.0 (Hopper), 10.0 (B200 / GB200), " + "10.3 (B300 / GB300) or 12.0 (RTX PRO 6000 Blackwell); this " + f"device reports {found}." + ) + if not platform._prepare_flash_attention_for_blackwell(): + raise RuntimeError( + "hybrid_window_attn_h3 requires FlashAttention for its dense legs" + ) + try: + from sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3 import ( # noqa: F401 + HybridWindowAttentionH3Backend, + ) + + return "sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3.HybridWindowAttentionH3Backend" + except Exception as e: + logger.error("Failed to import hybrid_window_attn_h3 backend: %s", str(e)) + raise ImportError( + "hybrid_window_attn_h3 needs FlashAttention and Triton." + ) from e + + class _CubeSparseAttentionBackendResolver(_CudaAttentionBackendResolver): backend = AttentionBackendEnum.CUBE_SPARSE_ATTN @@ -485,6 +531,7 @@ _CUDA_ATTENTION_BACKEND_RESOLVERS = { _SpargeAttentionBackendResolver, _VideoSparseAttentionBackendResolver, _VideoSparseAttentionH3BackendResolver, + _HybridWindowAttentionH3BackendResolver, _CubeSparseAttentionBackendResolver, _SparseVideoGen2AttentionBackendResolver, _SolAttnBackendResolver, @@ -652,7 +699,9 @@ class CudaPlatformBase(Platform): @classmethod def _prepare_flash_attention_for_blackwell(cls) -> bool: - if not cls.is_blackwell(): + # the FA4 CuTe package ships an sm120 forward kernel; the default FA backend + # still resolves to SDPA on SM120 before reaching this + if not (cls.is_blackwell() or cls.is_sm120()): return True try: diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 2b30835a2..129e92c5d 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -36,6 +36,7 @@ class AttentionBackendEnum(enum.Enum): SPARGE_ATTN = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() VIDEO_SPARSE_ATTN_H3 = enum.auto() + HYBRID_WINDOW_ATTN_H3 = enum.auto() SPARSE_VIDEO_GEN_2_ATTN = enum.auto() VMOBA_ATTN = enum.auto() AITER = enum.auto() @@ -59,6 +60,7 @@ class AttentionBackendEnum(enum.Enum): AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3, + AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3, AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SLA_ATTN, diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index 408332cc0..3daf4cc0d 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -46,6 +46,10 @@ BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = { "overlay_repo_id": "kevin-mi/FastH3-4step-Preview-overlay", "overlay_revision": "f769cb8001dae335089de7b250364335bc7cb183", }, + "OpenVDN/vdn-minimax-h3": { + "overlay_repo_id": "kevin-mi/VDN-H3-overlay", + "overlay_revision": "0ad315a05b914c4003af4d26152d288c2506a609", + }, } diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 5da7490cc..ace75c918 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -688,6 +688,47 @@ MINIMAX_H3_FOUR_GPU_H100_CASES = [ run_models_api_check=False, run_t2v_input_reference_check=False, ), + DiffusionTestCase( + "vdn_h3_t2va_4gpu_h100", + DiffusionServerArgs( + model_path="OpenVDN/vdn-minimax-h3", + modality="video", + num_gpus=4, + extras=[ + "--attention-backend", + "hybrid_window_attn_h3", + "--enable-torch-compile", + "false", + ], + ), + DiffusionSamplingParams( + prompt=( + "A curious raccoon peers through a vibrant field of yellow " + "sunflowers, its eyes wide with interest." + ), + output_size="1344x768", + seconds=5, + output_format="mp4", + expect_audio_output=True, + num_outputs_per_prompt=1, + extras={ + "task": "t2va", + "conditions": [], + "target": { + "short_edge": 768, + "aspect_ratio": "16:9", + "duration_seconds": 5.0, + }, + "num_inference_steps": 9, + "seed": 42, + }, + ), + run_perf_check=False, + run_consistency_check=False, + run_component_accuracy_check=False, + run_models_api_check=False, + run_t2v_input_reference_check=False, + ), DiffusionTestCase( "fasth3_t2va_vsa_4gpu_h100", DiffusionServerArgs( @@ -1428,6 +1469,7 @@ STANDALONE_FILES = { "../single_test_file/test_dp_serving_2_gpu.py", "../single_test_file/test_pynccl_a2a_capture_2_gpu.py", "../single_test_file/test_usp_replicated_parity_2_gpu.py", + "../single_test_file/test_vdn_ulysses_exchange_2_gpu.py", ], } @@ -1473,6 +1515,8 @@ STANDALONE_FILE_EST_TIMES = { "../single_test_file/test_pynccl_a2a_capture_2_gpu.py": 180.0, # two SDPA parity checks on 128+6 rows "../single_test_file/test_usp_replicated_parity_2_gpu.py": 180.0, + # no model load; two small all-to-alls + "../single_test_file/test_vdn_ulysses_exchange_2_gpu.py": 60.0, }, } diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_vdn_ulysses_exchange_2_gpu.py b/python/sglang/multimodal_gen/test/single_test_file/test_vdn_ulysses_exchange_2_gpu.py new file mode 100644 index 000000000..9fc78ba2f --- /dev/null +++ b/python/sglang/multimodal_gen/test/single_test_file/test_vdn_ulysses_exchange_2_gpu.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 Ulysses exchange: the field-major row->head all-to-all and its +inverse plus the head merge, checked against plain slicing on 2 GPUs.""" + +import torch +import torch.distributed as dist + +from sglang.test.test_utils import run_distributed_test + + +def _check(rank: int) -> None: + from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import ( + _vdn_a2a_heads_to_rows, + _vdn_a2a_rows_to_heads, + _vdn_merge_heads, + ) + + world = dist.get_world_size() + device = torch.device("cuda", rank) + heads, head_dim, local_rows = 6, 32, 24 + local_heads = heads // world + seq = local_rows * world + # every rank builds the same global tensors, so a shard is checked by slicing + g = torch.Generator(device="cpu").manual_seed(3) + qkv = torch.randn(seq, 3 * heads * head_dim, generator=g).to(device, torch.bfloat16) + q = qkv.view(seq, 3, heads, head_dim)[:, 0] # strided, as the split qkv is + scalars = torch.randn(seq, heads, 2, generator=g).to(device, torch.bfloat16) + rows = slice(rank * local_rows, (rank + 1) * local_rows) + mine = slice(rank * local_heads, (rank + 1) * local_heads) + with torch.inference_mode(): + for name, field in (("q", q), ("scalars", scalars)): + work, recv = _vdn_a2a_rows_to_heads( + field[rows], ulysses_ws=world, role=name, process_group=dist.group.WORLD + ) + work.wait() + assert recv.is_contiguous() + assert torch.equal(recv, field[:, mine]) + # inverse: this rank's heads for every row -> row shard, every head + work, back = _vdn_a2a_heads_to_rows( + q[:, mine].contiguous() * 2, + ulysses_ws=world, + role="out", + process_group=dist.group.WORLD, + ) + work.wait() + assert torch.equal(_vdn_merge_heads(back), q[rows] * 2) + torch.cuda.synchronize() + + +def test_exchange_round_trip_two_ranks() -> None: + run_distributed_test(_check, world_size=2) + + +if __name__ == "__main__": + test_exchange_round_trip_two_ranks() diff --git a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py index 1964cde61..024761f4b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py +++ b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py @@ -107,6 +107,51 @@ class TestCudaAttentionBackendSelection(unittest.TestCase): fake_flash_attn.set_fa_ver.assert_called_once_with(4) + def test_hybrid_window_h3_on_sm120_uses_fa4(self): + FakeCudaPlatform.is_sm120_device = True + FakeCudaPlatform.device_capability = DeviceCapability(12, 0) + fa_module = "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn" + fake_flash_attn = ModuleType(fa_module) + fake_flash_attn.set_fa_ver = Mock() + hybrid_module = "sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3" + fake_hybrid = ModuleType(hybrid_module) + fake_hybrid.HybridWindowAttentionH3Backend = object + + with patch.dict( + "sys.modules", {fa_module: fake_flash_attn, hybrid_module: fake_hybrid} + ): + self.assertEqual( + self.resolve(AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3), + f"{hybrid_module}.HybridWindowAttentionH3Backend", + ) + + fake_flash_attn.set_fa_ver.assert_called_once_with(4) + + def test_hybrid_window_h3_on_ampere_keeps_fa3(self): + FakeCudaPlatform.device_capability = DeviceCapability(8, 0) + fa_module = "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn" + fake_flash_attn = ModuleType(fa_module) + fake_flash_attn.set_fa_ver = Mock() + hybrid_module = "sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3" + fake_hybrid = ModuleType(hybrid_module) + fake_hybrid.HybridWindowAttentionH3Backend = object + + with patch.dict( + "sys.modules", {fa_module: fake_flash_attn, hybrid_module: fake_hybrid} + ): + self.assertEqual( + self.resolve(AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3), + f"{hybrid_module}.HybridWindowAttentionH3Backend", + ) + + fake_flash_attn.set_fa_ver.assert_not_called() + + def test_hybrid_window_h3_rejects_unsupported_capability(self): + FakeCudaPlatform.device_capability = DeviceCapability(7, 5) + + with self.assertRaisesRegex(ValueError, "12.0"): + self.resolve(AttentionBackendEnum.HYBRID_WINDOW_ATTN_H3) + def test_default_backend_uses_torch_sdpa_on_sm120(self): FakeCudaPlatform.is_sm120_device = True diff --git a/python/sglang/multimodal_gen/test/unit/test_hybrid_window_h3_attention.py b/python/sglang/multimodal_gen/test/unit/test_hybrid_window_h3_attention.py new file mode 100644 index 000000000..75a23c158 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_hybrid_window_h3_attention.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +"""hybrid_window_attn_h3 must reproduce a masked dense softmax with exactly the VDN +mask on a ragged packed layout, and dense attention once the window covers the clip.""" + +from __future__ import annotations + +import math +import sys + +import pytest +import torch + +from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import ( + VDNHybridAttentionArchConfig, +) +from sglang.multimodal_gen.runtime.layers.attention.backends.hybrid_window_attn_h3 import ( + HybridWindowAttentionH3Impl, + HybridWindowAttentionH3MetadataBuilder, + window_mask_frames, + window_mask_reference, +) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import VDNH3Layout +from sglang.multimodal_gen.runtime.platforms import current_platform + +# torch.cuda.is_available() is also True under ROCm, but the backend lives on the +# CUDA platform only (RocmPlatform rejects hybrid_window_attn_h3 outright and has +# no _prepare_flash_attention_for_blackwell), so gate on the platform itself. +requires_cuda = pytest.mark.skipif( + not current_platform.is_cuda(), + reason="hybrid_window_attn_h3 kernels need NVIDIA CUDA", +) + +# ragged on purpose: 70 and 100 are not tile multiples, 12 frames is not a chunk multiple +TEXT_LEN = 70 +AUDIO_ROWS = 100 +NUM_FRAMES = 12 +FRAME_H, FRAME_W = 6, 8 +TOKENS_PER_FRAME = FRAME_H * FRAME_W +HEADS = 4 +HEAD_DIM = 128 + + +def _layout() -> VDNH3Layout: + video_start = TEXT_LEN + AUDIO_ROWS + used = video_start + NUM_FRAMES * TOKENS_PER_FRAME + seq_len = (used + 63) // 64 * 64 + return VDNH3Layout( + seq_len=seq_len, + used=used, + text_len=TEXT_LEN, + video_start=video_start, + num_frames=NUM_FRAMES, + tokens_per_frame=TOKENS_PER_FRAME, + frame_height=FRAME_H, + frame_width=FRAME_W, + ) + + +def _hybrid(**overrides) -> VDNHybridAttentionArchConfig: + kwargs = dict(chunk=5, radius=1, anchor_frames="both") + kwargs.update(overrides) + return VDNHybridAttentionArchConfig(**kwargs) + + +def _qkv(device, seed: int = 7): + layout = _layout() + generator = torch.Generator(device="cpu").manual_seed(seed) + tensors = [ + torch.randn( + (layout.seq_len, HEADS, HEAD_DIM), generator=generator, dtype=torch.float32 + ).to(device=device, dtype=torch.bfloat16) + for _ in range(3) + ] + return layout, tensors + + +def _masked_reference(q, k, v, mask: torch.Tensor, used: int) -> torch.Tensor: + qf = q[:used].float().permute(1, 0, 2) + kf = k[:used].float().permute(1, 0, 2) + vf = v[:used].float().permute(1, 0, 2) + scores = qf @ kf.transpose(-2, -1) / math.sqrt(HEAD_DIM) + scores = scores.masked_fill(~mask[None], float("-inf")) + return (torch.softmax(scores, dim=-1) @ vf).permute(1, 0, 2) + + +def _prepare_flash_attention() -> None: + # the platform resolver runs this before the first forward; a direct impl must too + from sglang.multimodal_gen.runtime.platforms import current_platform + + current_platform._prepare_flash_attention_for_blackwell() + + +def _impl() -> HybridWindowAttentionH3Impl: + _prepare_flash_attention() + impl = HybridWindowAttentionH3Impl( + num_heads=HEADS, + head_size=HEAD_DIM, + causal=False, + softmax_scale=HEAD_DIM**-0.5, + num_kv_heads=HEADS, + prefix="blocks.3.attn", + ) + assert impl.layer_idx == 3 + return impl + + +def _run(impl, meta, layout, q, k, v, gate=None): + cu = torch.tensor( + [0, layout.used, layout.seq_len], dtype=torch.int32, device=q.device + ) + return impl.forward_varlen( + q, + k, + v, + cu_seqlens=cu, + max_seqlen=layout.used, + cu_seqlens_host=(0, layout.used, layout.seq_len), + attn_metadata=meta, + softmax_gate=gate, + ) + + +def test_window_bounds_and_anchor_frames() -> None: + hybrid = _hybrid() + bounds, dense_rows, dense_cols = window_mask_frames(hybrid, NUM_FRAMES) + # chunk 5, radius 1: frame 7 (chunk 1) sees chunks 0..2 = frames 0..14 -> clamped 11 + assert bounds[7] == (0, 11) + assert bounds[0] == (0, 9) + assert bounds[11] == (5, 11) + assert dense_rows == {0, NUM_FRAMES - 1} and dense_cols == {0, NUM_FRAMES - 1} + assert not hybrid.full_cover(NUM_FRAMES) + assert _hybrid(radius=NUM_FRAMES).full_cover(NUM_FRAMES) + # 102 = 20 * 5 + 2: the last chunk is short but still a whole chunk + raw = hybrid.window_bounds(102) + assert raw[100] == raw[101] == (95, 109) + assert raw[99] == (90, 104) + + +def test_mask_reference_partition_is_exact() -> None: + """Every (video q, video k) pair is in the softmax window or in the + linear branch's complement exactly once; anchors are absent from the + branch.""" + layout = _layout() + hybrid = _hybrid() + mask = window_mask_reference(hybrid, layout, torch.device("cpu")) + # globals dense both ways + assert mask[: layout.video_start].all() and mask[:, : layout.video_start].all() + assert mask[layout.video_end : layout.used].all() + bounds = hybrid.window_bounds(NUM_FRAMES) + vs, tpf = layout.video_start, TOKENS_PER_FRAME + for qf in range(NUM_FRAMES): + row = mask[vs + qf * tpf, vs : layout.video_end].view(NUM_FRAMES, tpf) + frame_kept = row.all(dim=1) + assert (frame_kept == row.any(dim=1)).all() # whole frames + if qf in (0, NUM_FRAMES - 1): + assert frame_kept.all() + continue + lo, hi = max(bounds[qf][0], 0), min(bounds[qf][1], NUM_FRAMES - 1) + expected_softmax = {f for f in range(lo, hi + 1)} | {0, NUM_FRAMES - 1} + # the branch covers the inner frames 1..F-2 outside the window + expected_linear = {f for f in range(1, NUM_FRAMES - 1) if f < lo or f > hi} + kept = {f for f in range(NUM_FRAMES) if frame_kept[f]} + assert kept == expected_softmax + assert kept.isdisjoint(expected_linear) + assert kept | expected_linear == set(range(NUM_FRAMES)) + + +@requires_cuda +def test_window_matches_masked_dense() -> None: + device = torch.device("cuda") + layout, (q, k, v) = _qkv(device) + hybrid = _hybrid() + meta = HybridWindowAttentionH3MetadataBuilder().build( + layout=layout, hybrid=hybrid, device=device + ) + assert not meta.full_cover + out = _run(_impl(), meta, layout, q, k, v) + mask = window_mask_reference(hybrid, layout, device) + reference = _masked_reference(q, k, v, mask, layout.used) + diff = (out[: layout.used].float() - reference).abs().max().item() + assert diff < 2e-2, f"window vs masked dense max diff {diff}" + assert torch.all(out[layout.used :] == 0) + + +@requires_cuda +def test_full_cover_matches_dense() -> None: + device = torch.device("cuda") + layout, (q, k, v) = _qkv(device, seed=11) + hybrid = _hybrid(radius=NUM_FRAMES) + meta = HybridWindowAttentionH3MetadataBuilder().build( + layout=layout, hybrid=hybrid, device=device + ) + assert meta.full_cover + out = _run(_impl(), meta, layout, q, k, v) + full = torch.ones(layout.used, layout.used, dtype=torch.bool, device=device) + reference = _masked_reference(q, k, v, full, layout.used) + diff = (out[: layout.used].float() - reference).abs().max().item() + assert diff < 2e-2, f"full cover vs dense max diff {diff}" + + +@requires_cuda +def test_decomposed_passes_are_arithmetic_neutral() -> None: + """Bounding the gathered K/V rows per pass splits the window into several + varlen calls without changing any query's kept set.""" + device = torch.device("cuda") + layout, (q, k, v) = _qkv(device, seed=9) + one = HybridWindowAttentionH3MetadataBuilder().build( + layout=layout, hybrid=_hybrid(), device=device + ) + many = HybridWindowAttentionH3MetadataBuilder().build( + layout=layout, hybrid=_hybrid(), device=device, max_gather_rows=1 + ) + assert len(one.decomposed.passes) == 1 and len(many.decomposed.passes) > 1 + impl = _impl() + a = _run(impl, one, layout, q, k, v).clone() + b = _run(impl, many, layout, q, k, v) + assert torch.equal(a, b) + + +@requires_cuda +def test_dense_fallback_off_the_dit_blocks() -> None: + """The token refiner resolves the same backend but runs plain dense FA.""" + device = torch.device("cuda") + layout, (q, k, v) = _qkv(device, seed=5) + _prepare_flash_attention() + impl = HybridWindowAttentionH3Impl( + num_heads=HEADS, + head_size=HEAD_DIM, + causal=False, + softmax_scale=HEAD_DIM**-0.5, + num_kv_heads=HEADS, + prefix="token_refiner.blocks.0.attn", + ) + assert impl.layer_idx is None + out = _run(impl, None, layout, q, k, v) + full = torch.ones(layout.used, layout.used, dtype=torch.bool, device=device) + reference = _masked_reference(q, k, v, full, layout.used) + assert (out[: layout.used].float() - reference).abs().max().item() < 2e-2 + + +@requires_cuda +def test_metadata_layout_mismatch_is_rejected() -> None: + device = torch.device("cuda") + layout, (q, k, v) = _qkv(device) + meta = HybridWindowAttentionH3MetadataBuilder().build( + layout=layout, hybrid=_hybrid(), device=device + ) + cu = torch.tensor( + [0, layout.used - 48, layout.seq_len], dtype=torch.int32, device=device + ) + with pytest.raises(ValueError, match="diverged"): + _impl().forward_varlen( + q, + k, + v, + cu_seqlens=cu, + max_seqlen=layout.used - 48, + cu_seqlens_host=(0, layout.used - 48, layout.seq_len), + attn_metadata=meta, + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py index c99ce719c..ab7393cf5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py @@ -238,7 +238,9 @@ def test_cache_dit_preservation_only_makes_first_gate_out_of_place(): block.norm1 = torch.nn.Identity() block.norm2 = torch.nn.Identity() block.attn = _KwargIdentity() + block.attn.qkv_proj = SimpleNamespace(quant_method=UnquantizedLinearMethod()) block.mlp = torch.nn.Identity() + block.mlp.fc1 = SimpleNamespace(quant_method=UnquantizedLinearMethod()) gate_modes = [] def fake_gate(residual, _gate, _other, _indices, *, dtype, allow_inplace=True): diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vdn.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vdn.py new file mode 100644 index 000000000..6a330c510 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vdn.py @@ -0,0 +1,703 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VDN-H3 (hybrid window-softmax + Video Delta linear attention MiniMax-H3): +registration, admission, and the linear branch's arithmetic against +step-by-step references (no weights, CPU + small CUDA shapes).""" + +from __future__ import annotations + +import math +import re +import sys +from types import SimpleNamespace + +import pytest +import torch + +from sglang.multimodal_gen.configs.models.dits.minimax_h3 import ( + MiniMaxH3DiTArchConfig, + MiniMaxH3DiTConfig, + VDNHybridAttentionArchConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( + MiniMaxH3PipelineConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3_vdn import ( + VDNH3PipelineConfig, +) +from sglang.multimodal_gen.configs.sample.minimax_h3_vdn import VDNH3SamplingParams +from sglang.multimodal_gen.registry import ( + get_model_info, + get_non_diffusers_pipeline_name, +) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import ( + TEXT_STATE_SCALE, + MiniMaxH3VDNLinearBranch, + VDNH3Layout, + delta_factor_apply, + frame_statistics, + gather_linear_state, + run_scans, + vdn_h3_layout_from_packed, +) +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, + current_platform, +) + +VDN_MODEL_ID = "OpenVDN/vdn-minimax-h3" +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +# admission resolves hybrid_window_attn_h3, which only the CUDA platform registers +requires_cuda_backend = pytest.mark.skipif( + not current_platform.is_cuda(), + reason="hybrid_window_attn_h3 admission needs NVIDIA CUDA", +) + + +# -------------------------------------------------------------------------- +# registration and admission +# -------------------------------------------------------------------------- + + +def test_registry_resolves_vdn_h3_configs() -> None: + info = get_model_info(VDN_MODEL_ID) + assert info.sampling_param_cls is VDNH3SamplingParams + assert info.pipeline_config_cls is VDNH3PipelineConfig + assert get_non_diffusers_pipeline_name(VDN_MODEL_ID) == "VDNH3Pipeline" + assert get_non_diffusers_pipeline_name("/models/OpenVDN/vdn-minimax-h3") == ( + "VDNH3Pipeline" + ) + # the base H3 detector must not swallow the VDN id (it contains "minimax-h3") + base = get_model_info("MiniMaxAI/MiniMax-H3") + assert base.pipeline_config_cls is MiniMaxH3PipelineConfig + + +def test_vdn_h3_sampling_defaults_and_rejections() -> None: + params = VDNH3SamplingParams(prompt="p") + assert params.num_inference_steps == 9 # 8 NFE + with pytest.raises(ValueError, match="exactly nine sigma grid points"): + VDNH3SamplingParams(prompt="p", num_inference_steps=8) + fl2va = VDNH3SamplingParams( + prompt="p", + task="fl2va", + conditions=[ + {"type": "image", "uri": "x.png", "role": "keyframe", "frame_index": 0} + ], + target={"short_edge": 768, "aspect_ratio": "auto", "duration_seconds": 5.0}, + ) + assert fl2va.task == "fl2va" + with pytest.raises(ValueError, match="ref2va was not trained"): + VDNH3SamplingParams( + prompt="p", + task="ref2va", + conditions=[{"type": "image", "uri": "x.png", "role": "reference"}], + target={"short_edge": 768, "aspect_ratio": "auto", "duration_seconds": 5.0}, + ) + + +def _server_args(**overrides) -> SimpleNamespace: + args = dict( + model_variant=None, + attention_backend=None, + component_attention_backends={}, + attention_backend_config=None, + ring_degree=1, + enable_torch_compile=False, + enable_breakable_cuda_graph=False, + quantization=None, + ) + args.update(overrides) + ns = SimpleNamespace(**args) + ns.resolve_component_attention_backend = lambda name: ( + ( + AttentionBackendEnum[str(ns.component_attention_backends[name]).upper()] + if name in ns.component_attention_backends + else None + ), + None, + ) + return ns + + +@requires_cuda_backend +def test_vdn_h3_pipeline_config_rejections() -> None: + config = VDNH3PipelineConfig() + with pytest.raises(ValueError, match="--model-variant does not apply"): + config.validate_server_args(_server_args(model_variant="ref2va")) + with pytest.raises( + ValueError, match="requires --attention-backend hybrid_window_attn_h3" + ): + config.validate_server_args(_server_args(attention_backend="fa")) + with pytest.raises(ValueError, match="ring-degree"): + config.validate_server_args(_server_args(ring_degree=2)) + with pytest.raises(ValueError, match="torch.compile"): + config.validate_server_args(_server_args(enable_torch_compile=True)) + with pytest.raises(ValueError, match="breakable CUDA graph"): + config.validate_server_args(_server_args(enable_breakable_cuda_graph=True)) + with pytest.raises(ValueError, match="no.*audited high-quality deployment"): + config.validate_quality_deployment(server_args=None) + args = _server_args() + config.validate_server_args(args) + assert args.attention_backend == "hybrid_window_attn_h3" + + +@requires_cuda_backend +def test_vdn_h3_quantization_defaults_to_mxfp8_on_blackwell(monkeypatch) -> None: + """Online MXFP8 is the default on SM100+ (SM120 included) and what `fp8` + maps to there; `bf16` opts out; before SM100 the block-scaled GEMM does + not exist, so an unset flag stays bf16 and `fp8` stays the per-channel + path.""" + from sglang.multimodal_gen.configs.pipeline_configs import minimax_h3_vdn as module + + def resolved( + quantization: str | None, blackwell: bool, sm120: bool = False + ) -> str | None: + monkeypatch.setattr(module.current_platform, "is_blackwell", lambda: blackwell) + monkeypatch.setattr(module.current_platform, "is_sm120", lambda: sm120) + args = _server_args(quantization=quantization) + VDNH3PipelineConfig().validate_server_args(args) + return args.quantization + + assert resolved(None, True) == "mxfp8" + assert resolved("fp8", True) == "mxfp8" + assert resolved("bf16", True) is None + assert resolved(None, False) is None + assert resolved("fp8", False) == "fp8" + assert resolved(None, False, sm120=True) == "mxfp8" + assert resolved("fp8", False, sm120=True) == "mxfp8" + assert resolved("bf16", False, sm120=True) is None + + +def test_hybrid_arch_config_from_transform_config_and_mapping() -> None: + transform = { + "anchor_frames": "both", + "enable_softmax_gate": True, + "linear_attention": { + "a_fp32": True, + "bridge": "alpha", + "delta_rule": "vdn_solve", + "enable_text_state": True, + "linear_head_dim": 128, + "short_conv": {"targets": ["k", "v"]}, + }, + "softmax_attention": {"chunk": 5, "radius": 1}, + } + dit = MiniMaxH3DiTConfig() + dit.update_model_arch({"hybrid_attention": transform, "num_layers": 2}) + hybrid = dit.arch_config.hybrid_attention + assert isinstance(hybrid, VDNHybridAttentionArchConfig) + assert hybrid.short_conv == ("k", "v") and hybrid.chunk == 5 + assert MiniMaxH3DiTConfig().arch_config.hybrid_attention is None + with pytest.raises(ValueError, match="delta_rule"): + VDNHybridAttentionArchConfig(delta_rule="bogus") + + mapping = MiniMaxH3DiTArchConfig().param_names_mapping + for source, expected in ( + ( + "transformer_blocks.7.attn.linear_attention.alpha.A_log", + "blocks.7.attn.hybrid.linear_attention.alpha.A_log", + ), + ( + "transformer_blocks.7.attn.softmax_gate.up.bias", + "blocks.7.attn.hybrid.softmax_gate.up.bias", + ), + ( + "transformer_blocks.7.attn.to_out_linear.weight", + "blocks.7.attn.hybrid.to_out_linear.weight", + ), + ): + targets = [ + re.sub(pattern, target if isinstance(target, str) else target[0], source) + for pattern, target in mapping.items() + if re.match(pattern, source) + ] + assert targets == [expected], (source, targets) + + +def test_layout_from_packed_t2va() -> None: + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + ) + + packed = minimax_h3_packed_sequence( + text_len=70, + latent_t=12, + latent_h=12, + latent_w=16, + audio_t=50, + include_keyframe_cond=False, + ) + layout = vdn_h3_layout_from_packed(packed, latent_t=12, latent_h=12, latent_w=16) + assert layout.text_len == 70 + assert layout.video_start == 70 + 100 + assert layout.tokens_per_frame == 48 and layout.frame_size == (6, 8) + assert layout.used == 70 + 100 + 12 * 48 + assert layout.seq_len == int(packed["seq_len"]) and layout.seq_len % 64 == 0 + assert layout.global_ranges == [(0, 170)] + + +def test_layout_from_packed_fl2va_keeps_keyframe_rows_global() -> None: + """Keyframe rows must land in the global ranges, not in the video span + the linear branch scans.""" + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + ) + + packed = minimax_h3_packed_sequence( + text_len=70, + latent_t=12, + latent_h=12, + latent_w=16, + audio_t=50, + include_keyframe_cond=True, + keyframe_frame_indices=[0, -1], + frame_count=45, + ) + layout = vdn_h3_layout_from_packed(packed, latent_t=12, latent_h=12, latent_w=16) + assert layout.text_len == 70 + assert layout.video_start == 70 + 2 * 48 + 100 + assert layout.used == layout.video_end == layout.video_start + 12 * 48 + assert layout.global_ranges == [(0, layout.video_start)] + + +# -------------------------------------------------------------------------- +# the linear branch arithmetic +# -------------------------------------------------------------------------- + +FRAMES, HEADS, TOKENS, HEAD_DIM = 6, 2, 16, 32 + + +def _random_stats(device, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + k = torch.nn.functional.normalize( + torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g), dim=-1 + ).to(device) + v = torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g).to(device) + beta = torch.rand(FRAMES, HEADS, TOKENS, generator=g).to(device) + alpha = torch.rand(FRAMES, HEADS, HEAD_DIM, generator=g).to(device) * 0.5 + 0.5 + return k, v, beta, alpha + + +def test_frame_statistics_and_delta_rule_match_dense_algebra() -> None: + k, v, beta, alpha = _random_stats("cpu") + A, B = frame_statistics(k, v, beta, a_fp32=True) + A_ref = torch.einsum("fhsk,fhs,fhsl->fhkl", k, beta, k) + B_ref = torch.einsum("fhsv,fhs,fhsk->fhvk", v, beta, k) + assert torch.allclose(A, A_ref, atol=1e-4) and torch.allclose(B, B_ref, atol=1e-4) + transition, injection = delta_factor_apply( + "vdn_solve", alpha, A, B, tokens_per_frame=TOKENS + ) + inv = torch.linalg.inv(torch.eye(HEAD_DIM) + A) + assert torch.allclose(transition, alpha.unsqueeze(-1) * inv, atol=1e-4) + assert torch.allclose(injection, B @ inv, atol=1e-4) + + +def test_scans_match_step_reference_and_text_seed() -> None: + k, v, beta, alpha = _random_stats("cpu", seed=1) + A, B = frame_statistics(k, v, beta, a_fp32=True) + transition, injection = delta_factor_apply( + "vdn_solve", alpha, A, B, tokens_per_frame=TOKENS + ) + text_state = torch.randn(HEADS, HEAD_DIM, HEAD_DIM) + prefix, suffix = run_scans(transition, injection, text_state) + state = text_state.clone() + for f in range(FRAMES): + state = state @ transition[f] + injection[f] + assert torch.allclose(prefix[f], state, atol=1e-4) + state = text_state.clone() + for f in range(FRAMES - 1, -1, -1): + state = state @ transition[f] + injection[f] + assert torch.allclose(suffix[f], state, atol=1e-4) + + +@pytest.mark.parametrize("world", [1, 2, 5, 10]) +def test_frame_partial_sums_match_index_add(world: int) -> None: + """The Ulysses frame-mean partial sums (reshape-sum over whole frames plus + two edge rows sums, deterministic) equal the index_add formulation.""" + from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn_attention import ( + _vdn_frame_partial_sums, + ) + + frames, tpf, hidden = 12, 50, 64 + video_start = 20 + video_end = video_start + frames * tpf + seq = video_end + 30 # 650 rows, divisible by every ``world`` above + g = torch.Generator(device="cpu").manual_seed(0) + x = torch.randn(seq, hidden, generator=g).to(torch.bfloat16) + local = seq // world + total = torch.zeros(frames, hidden) + for rank in range(world): + total += _vdn_frame_partial_sums( + x[rank * local : (rank + 1) * local], + row_start=rank * local, + video_start=video_start, + video_end=video_end, + num_frames=frames, + tokens_per_frame=tpf, + ) + ref = x[video_start:video_end].float().view(frames, tpf, hidden).sum(1) + torch.testing.assert_close(total, ref, rtol=1e-5, atol=1e-3) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("anchor_frames", ["both", "none"]) +@pytest.mark.parametrize("reference", ["frame_chain_scans", "eager_kernels"]) +def test_branch_forward_matches_reference( + anchor_frames: str, reference: str, monkeypatch +) -> None: + """The shipped branch (boundary scans, fused Triton kernels) against the + same module with the plain frame-chain scans, and against the eager + kernel chain; both across the anchor-frame shift of the chunk grid.""" + from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as module + + device = torch.device("cuda") + hidden, heads, head_dim = 64, 4, 32 + num_frames, fh, fw = 13, 4, 6 + tpf = fh * fw + hybrid = VDNHybridAttentionArchConfig( + chunk=3, radius=1, anchor_frames=anchor_frames, linear_head_dim=head_dim + ) + branch = _branch(hybrid, heads, hidden, head_dim, seed=0).to(device) + layout = VDNH3Layout( + seq_len=64 * 6, + used=10 + num_frames * tpf, + text_len=10, + video_start=10, + num_frames=num_frames, + tokens_per_frame=tpf, + frame_height=fh, + frame_width=fw, + ) + g = torch.Generator(device="cpu").manual_seed(1) + V = num_frames * tpf + q, k, v = ( + torch.randn(V, heads, head_dim, generator=g).to(device, torch.bfloat16) + for _ in range(3) + ) + tk, tv = ( + torch.randn(10, heads, head_dim, generator=g).to(device, torch.bfloat16) + for _ in range(2) + ) + x = torch.randn(V, hidden, generator=g).to(device, torch.bfloat16) + tx = torch.randn(10, hidden, generator=g).to(device, torch.bfloat16) + kwargs = dict( + q_raw=q, + k_raw=k, + v_raw=v, + beta=branch.beta(x), + gate=branch.output_gate(x), + frame_mean=x.view(num_frames, tpf, hidden).mean(1, dtype=torch.float32), + layout=layout, + text_k_raw=tk, + text_v_raw=tv, + text_beta=branch.beta(tx), + ) + fast = branch(**kwargs) + if reference == "frame_chain_scans": + monkeypatch.setattr( + module, + "run_boundary_scans", + lambda t, i, ts, *, chunk, frame_offset=0: run_scans(t, i, ts), + ) + else: + branch.fused_kernels = False + expected = branch(**kwargs) + assert expected.abs().sum() > 0 + # fused kernels skip the eager chain's bf16 roundings; the scans re-associate fp32 + tolerance = {"frame_chain_scans": 5e-3, "eager_kernels": 2e-2}[reference] + rel = (fast.float() - expected.float()).norm() / expected.float().norm() + assert rel < tolerance, rel + + +def test_gather_is_the_exact_window_complement() -> None: + """With alpha = 1 and one-hot frame indicators, the gathered state must + be exactly the indicator of the frames outside the window.""" + num_frames = 9 + hybrid = VDNHybridAttentionArchConfig(chunk=3, radius=1, anchor_frames="none") + bounds = hybrid.window_bounds(num_frames) + # injection[f] = one-hot(f) laid along dv; transition = identity + eye = torch.eye(num_frames) + injection = ( + eye.view(num_frames, 1, num_frames, 1) + .expand(num_frames, 1, num_frames, 1) + .clone() + ) + transition = torch.eye(1).view(1, 1, 1, 1).expand(num_frames, 1, 1, 1).clone() + prefix, suffix = run_scans(transition, injection, None) + alpha = torch.ones(num_frames, 1, 1) + gathered = gather_linear_state( + prefix, + suffix, + alpha, + bounds, + bridge="alpha", + text_state=None, + out_dtype=torch.float32, + ) + for t in range(num_frames): + lo, hi = max(bounds[t][0], 0), min(bounds[t][1], num_frames - 1) + expected = torch.tensor( + [1.0 if (f < lo or f > hi) else 0.0 for f in range(num_frames)] + ) + assert torch.equal(gathered[t, 0, :, 0], expected), (t, gathered[t, 0, :, 0]) + + +def test_gather_text_state_decays_over_skipped_frames() -> None: + """A clip-end frame reads the text state decayed by prod alpha over + exactly the frames between the boundary and t (VDN's bridge indices).""" + num_frames = 4 + bounds = [ + (t, t) for t in range(num_frames) + ] # radius 0: complement = everything else + prefix = torch.zeros(num_frames, 1, 1, 1) + suffix = torch.zeros(num_frames, 1, 1, 1) + alpha = torch.tensor([0.5, 0.25, 0.5, 0.5]).view(num_frames, 1, 1) + text_state = torch.ones(1, 1, 1) + out = gather_linear_state( + prefix, + suffix, + alpha, + bounds, + bridge="alpha", + text_state=text_state, + out_dtype=torch.float32, + ).view(num_frames) + # frame 0 reads the text state through alpha[0]; frame 3 through alpha[3] + assert math.isclose(out[0].item(), 0.5, rel_tol=1e-6) + assert math.isclose(out[3].item(), 0.5, rel_tol=1e-6) + assert out[1].item() == 0.0, "both neighbours in range and zero" + + +def test_temporal_shift_features_match_conv1d() -> None: + from sglang.multimodal_gen.runtime.models.dits.minimax_h3_vdn import _temporal_shift + + x = torch.randn(7, 5, 6) # [F, S, C] + w = torch.randn(6, 5) + got = _temporal_shift(x, w) + ref = ( + torch.nn.functional.conv1d( + x.permute(1, 2, 0).reshape(5, 6, 7), w.view(6, 1, 5), padding=2, groups=6 + ) + .reshape(5, 6, 7) + .permute(2, 0, 1) + ) + assert torch.allclose(got, ref, atol=1e-5) + + +def _branch( + hybrid: VDNHybridAttentionArchConfig, + heads: int, + hidden: int, + head_dim: int, + seed: int, +): + from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, + ) + from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + ensure_distributed_env_defaults, + ) + + if not model_parallel_is_initialized(): + ensure_distributed_env_defaults() + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + arch = MiniMaxH3DiTArchConfig( + num_attention_heads=heads, attention_head_dim=head_dim, hidden_size=hidden + ) + branch = MiniMaxH3VDNLinearBranch(arch, hybrid, local_heads=heads) + g = torch.Generator(device="cpu").manual_seed(seed) + with torch.no_grad(): + for name, p in branch.named_parameters(): + if name.endswith("A_log"): + p.copy_(torch.log(torch.empty_like(p).uniform_(1, 16, generator=g))) + elif name.endswith("dt_bias"): + p.copy_(torch.randn(p.shape, generator=g) - 3) + elif name.endswith("norm.weight"): + p.copy_(torch.ones_like(p)) + else: + p.copy_( + torch.randn(p.shape, generator=g, dtype=torch.float32).to(p.dtype) + * 0.1 + ) + return branch + + +@requires_cuda +def test_branch_head_slice_equals_full_run() -> None: + """The Ulysses contract: the branch is per-head independent given + (beta, gate, alpha), so a head-sliced run equals the full run's slice.""" + device = torch.device("cuda") + hidden, heads, head_dim = 64, 4, 32 + num_frames, fh, fw = 7, 4, 6 + tpf = fh * fw + hybrid = VDNHybridAttentionArchConfig( + chunk=2, radius=1, anchor_frames="both", linear_head_dim=head_dim + ) + branch = _branch(hybrid, heads, hidden, head_dim, seed=0).to(device) + layout = VDNH3Layout( + seq_len=64 * 4, + used=10 + num_frames * tpf, + text_len=10, + video_start=10, + num_frames=num_frames, + tokens_per_frame=tpf, + frame_height=fh, + frame_width=fw, + ) + g = torch.Generator(device="cpu").manual_seed(1) + V = num_frames * tpf + q, k, v = ( + torch.randn(V, heads, head_dim, generator=g).to(device, torch.bfloat16) + for _ in range(3) + ) + tk, tv = ( + torch.randn(10, heads, head_dim, generator=g).to(device, torch.bfloat16) + for _ in range(2) + ) + x = torch.randn(V, hidden, generator=g).to(device, torch.bfloat16) + tx = torch.randn(10, hidden, generator=g).to(device, torch.bfloat16) + beta, gate = branch.beta(x), branch.output_gate(x) + tbeta = branch.beta(tx) + frame_mean = x.view(num_frames, tpf, hidden).mean(1, dtype=torch.float32) + full = branch( + q_raw=q, + k_raw=k, + v_raw=v, + beta=beta, + gate=gate, + frame_mean=frame_mean, + layout=layout, + text_k_raw=tk, + text_v_raw=tv, + text_beta=tbeta, + ).view(V, heads, head_dim) + # anchors read zero + assert torch.all(full[:tpf] == 0) and torch.all(full[-tpf:] == 0) + assert full[tpf:-tpf].abs().sum() > 0 + + # the same module on a head range of the full sequence + hs = slice(1, 3) + part = branch( + q_raw=q[:, hs], + k_raw=k[:, hs], + v_raw=v[:, hs], + beta=beta[:, hs], + gate=gate[:, hs], + frame_mean=frame_mean, + layout=layout, + text_k_raw=tk[:, hs], + text_v_raw=tv[:, hs], + text_beta=tbeta[:, hs], + heads=hs, + ).view(V, 2, head_dim) + diff = (part.float() - full[:, hs].float()).abs().max().item() + assert diff < 2e-2, f"head slice vs full run max diff {diff}" + + +@requires_cuda +def test_branch_matches_eager_reference_algorithm() -> None: + """The module against a from-scratch spelling of VDN's _readout (no + skip_ends), including the text state seed.""" + device = torch.device("cuda") + hidden, heads, head_dim = 48, 2, 32 + num_frames, fh, fw = 5, 3, 4 + tpf = fh * fw + hybrid = VDNHybridAttentionArchConfig( + chunk=0, radius=1, anchor_frames="none", linear_head_dim=head_dim, short_conv=() + ) + branch = _branch(hybrid, heads, hidden, head_dim, seed=2).to(device) + layout = VDNH3Layout( + seq_len=256, + used=8 + num_frames * tpf, + text_len=8, + video_start=8, + num_frames=num_frames, + tokens_per_frame=tpf, + frame_height=fh, + frame_width=fw, + ) + g = torch.Generator(device="cpu").manual_seed(3) + V = num_frames * tpf + q, k, v = ( + torch.randn(V, heads, head_dim, generator=g).to(device, torch.bfloat16) + for _ in range(3) + ) + tk, tv = ( + torch.randn(8, heads, head_dim, generator=g).to(device, torch.bfloat16) + for _ in range(2) + ) + x = torch.randn(V, hidden, generator=g).to(device, torch.bfloat16) + tx = torch.randn(8, hidden, generator=g).to(device, torch.bfloat16) + beta, gate, tbeta = branch.beta(x), branch.output_gate(x), branch.beta(tx) + frame_mean = x.view(num_frames, tpf, hidden).mean(1, dtype=torch.float32) + got = branch( + q_raw=q, + k_raw=k, + v_raw=v, + beta=beta, + gate=gate, + frame_mean=frame_mean, + layout=layout, + text_k_raw=tk, + text_v_raw=tv, + text_beta=tbeta, + ) + + # reference in fp32 + def feat(t, l2): + y = torch.nn.functional.silu(t.float()) + return torch.nn.functional.normalize(y, dim=-1, eps=1e-6) if l2 else y + + qf, kf, vf = feat(q, True), feat(k, True), feat(v, False) + bounds = hybrid.window_bounds(num_frames) + kb = kf.view(num_frames, tpf, heads, head_dim).permute(0, 2, 1, 3) + vb = vf.view(num_frames, tpf, heads, head_dim).permute(0, 2, 1, 3) + bb = beta.float().view(num_frames, tpf, heads).permute(0, 2, 1) + A = torch.einsum("fhsk,fhs,fhsl->fhkl", kb, bb, kb) + B = torch.einsum("fhsv,fhs,fhsk->fhvk", vb, bb, kb) + inv = torch.linalg.inv(torch.eye(head_dim, device=device) + A) + alpha = branch.alpha(frame_mean) + trans = alpha.unsqueeze(-1) * inv + inj = B @ inv + # text state + tkf, tvf = feat(tk, True), feat(tv, False) + tkb = tkf.view(1, 8, heads, head_dim).permute(0, 2, 1, 3) + tvb = tvf.view(1, 8, heads, head_dim).permute(0, 2, 1, 3) + tbb = tbeta.float().view(1, 8, heads).permute(0, 2, 1) + tA = torch.einsum("fhsk,fhs,fhsl->fhkl", tkb, tbb, tkb)[0] + tB = torch.einsum("fhsv,fhs,fhsk->fhvk", tvb, tbb, tkb)[0] + text_state = TEXT_STATE_SCALE * ( + tB @ torch.linalg.inv(torch.eye(head_dim, device=device) + tA) + ) + prefix, suffix = [], [None] * num_frames + s = text_state.clone() + for f in range(num_frames): + s = s @ trans[f] + inj[f] + prefix.append(s) + s = text_state.clone() + for f in range(num_frames - 1, -1, -1): + s = s @ trans[f] + inj[f] + suffix[f] = s + outs = [] + for t in range(num_frames): + lo, hi = bounds[t] + left = prefix[lo - 1] if lo - 1 >= 0 else text_state + right = suffix[hi + 1] if hi + 1 < num_frames else text_state + a_before = torch.prod(alpha[max(lo, 0) : t + 1], dim=0) + a_after = torch.prod(alpha[t : min(hi, num_frames - 1) + 1], dim=0) + state = left * a_before.unsqueeze(1) + right * a_after.unsqueeze(1) + qt = qf.view(num_frames, tpf, heads, head_dim)[t] # [S, H, d] + ro = torch.einsum("shk,hvk->shv", qt, state) + ms = ro.pow(2).mean(-1, keepdim=True) + ro = ro * torch.rsqrt(ms + branch.norm.eps) * branch.norm.weight.float() + outs.append(ro) + ref = (torch.cat(outs) * gate.float()).reshape(V, heads * head_dim) + diff = (got.float() - ref).abs().max().item() + scale = ref.abs().max().item() + assert diff < 3e-2 * max(scale, 1.0), ( + f"branch vs reference max diff {diff} (scale {scale})" + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/multimodal_gen/test/unit/test_mxfp8_online_gemm.py b/python/sglang/multimodal_gen/test/unit/test_mxfp8_online_gemm.py new file mode 100644 index 000000000..9758d57d2 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_mxfp8_online_gemm.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Online MXFP8 GEMM path (``MXFP8Config()`` on a bf16 checkpoint): load-time +block quant, the prequantized (e4m3, swizzled scales) input from the fused +SwiGLU kernel, and the per-layer fallback to the per-channel fp8 path.""" + +import sys + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _init_parallel() -> None: + from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, + ) + from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + ensure_distributed_env_defaults, + ) + + if not model_parallel_is_initialized(): + ensure_distributed_env_defaults() + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + + +def _layer( + in_f: int, out_f: int, bias: bool, params_dtype: torch.dtype = torch.bfloat16 +): + from sglang.multimodal_gen.runtime.layers.linear import RowParallelLinear + from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config + + return RowParallelLinear( + in_f, + out_f, + bias=bias, + params_dtype=params_dtype, + quant_config=MXFP8Config(), + prefix="mlp.fc2", + ).to("cuda") + + +def test_mxfp8_linear_matches_bf16_and_accepts_prequantized() -> None: + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("cuBLASLt MXFP8 block scaling requires Blackwell or newer") + _init_parallel() + from sglang.kernels.ops.diffusion import silu_mul_mxfp8 + + in_f, out_f, rows = 512, 384, 200 + layer = _layer(in_f, out_f, bias=False) + g = torch.Generator(device="cpu").manual_seed(1) + weight = (torch.randn(out_f, in_f, generator=g) * 0.02).to("cuda", torch.bfloat16) + with torch.no_grad(): + layer.weight.copy_(weight) + layer.quant_method.process_weights_after_loading(layer) + assert layer.mxfp8 and layer.quant_method.accepts_mxfp8_input(layer) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight_scale.dtype == torch.float8_e8m0fnu + + x = torch.randn(rows, in_f, generator=g).to("cuda", torch.bfloat16) + out, _ = layer(x) + ref = x.float() @ weight.float().t() + rel = ((out.float() - ref).norm() / ref.norm()).item() + assert rel < 0.05, rel + + hidden = torch.randn(rows, 2 * in_f, generator=g).to("cuda", torch.bfloat16) + act = torch.nn.functional.silu(hidden[:, :in_f]) * hidden[:, in_f:] + out_tensor, _ = layer(act) + out_tuple, _ = layer(silu_mul_mxfp8(hidden)) + assert torch.equal(out_tensor, out_tuple) + + +def _assert_aligned_layer_falls_back(params_dtype: torch.dtype) -> None: + _init_parallel() + layer = _layer(512, 384, bias=False, params_dtype=params_dtype) + g = torch.Generator(device="cpu").manual_seed(1) + weight = (torch.randn(384, 512, generator=g) * 0.02).to("cuda", params_dtype) + with torch.no_grad(): + layer.weight.copy_(weight) + layer.quant_method.process_weights_after_loading(layer) + assert not layer.mxfp8 and not layer.quant_method.accepts_mxfp8_input(layer) + x = torch.randn(200, 512, generator=g).to("cuda", params_dtype) + out, _ = layer(x) + ref = x.float() @ weight.float().t() + rel = ((out.float() - ref).norm() / ref.norm()).item() + assert rel < 0.05, rel + + +def test_pre_blackwell_aligned_layer_falls_back_to_channelwise() -> None: + if torch.cuda.get_device_capability()[0] >= 10: + pytest.skip("requires a pre-Blackwell GPU") + _assert_aligned_layer_falls_back(torch.bfloat16) + + +def test_fp16_layer_falls_back_to_channelwise() -> None: + """The swizzled quantizer takes bf16 only; an fp16 layer must not fail at load.""" + _assert_aligned_layer_falls_back(torch.float16) + + +def test_unaligned_layer_falls_back_to_channelwise() -> None: + """The block-scaled GEMM needs K % 32 == 0; such a layer keeps the + per-channel fp8 path and still answers a forward. K is 16 rather than a + smaller odd size because the fallback's scaled GEMM still wants K % 16 == 0 + (ROCm rejects anything else outright).""" + _init_parallel() + layer = _layer(16, 128, bias=True) + layer.quant_method.process_weights_after_loading(layer) + assert not layer.mxfp8 and not layer.quant_method.accepts_mxfp8_input(layer) + out, _ = layer(torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)) + assert out.shape == (4, 128) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernel/diffusion/benchmark/bench_vdn_delta_factors.py b/test/registered/kernel/diffusion/benchmark/bench_vdn_delta_factors.py new file mode 100644 index 000000000..ee84cebe3 --- /dev/null +++ b/test/registered/kernel/diffusion/benchmark/bench_vdn_delta_factors.py @@ -0,0 +1,51 @@ +"""``vdn_delta_factors`` (fused inverse + products) vs the eager Cholesky chain it replaces.""" + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.diffusion import vdn_delta_factors +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +HEAD_DIM = 128 + + +def eager_delta_factors(A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor): + eye = torch.eye(HEAD_DIM, device=A.device, dtype=torch.float32).expand_as(A) + chol = torch.linalg.cholesky(A + eye) + linv = torch.linalg.solve_triangular(chol, eye, upper=False, left=True) + inv = linv.transpose(-1, -2) @ linv + return alpha.unsqueeze(-1) * inv, B @ inv + + +FN_MAP = {"jit": vdn_delta_factors, "eager": eager_delta_factors} + + +def _inputs(num: int): + g = torch.Generator(device="cuda").manual_seed(0) + k = torch.nn.functional.normalize( + torch.randn(num, 1008, HEAD_DIM, device="cuda", generator=g), dim=-1 + ) + v = torch.randn(num, 1008, HEAD_DIM, device="cuda", generator=g) + beta = torch.sigmoid(torch.randn(num, 1008, device="cuda", generator=g)) + A = (k * beta.unsqueeze(-1)).transpose(-1, -2) @ k + A = 0.5 * (A + A.transpose(-1, -2)) + B = (v * beta.unsqueeze(-1)).transpose(-1, -2) @ k + alpha = torch.rand(num, HEAD_DIM, device="cuda", generator=g) + return A.contiguous(), B.contiguous(), alpha.contiguous() + + +# 707 = 101 frames x 7 heads: the paper workload per rank (8 x B200, Ulysses 8) +@marker.parametrize("num_matrices", [64, 707], [64]) +@marker.benchmark("impl", ["jit", "eager"], unit="us") +def benchmark(num_matrices: int, impl: str): + A, B, alpha = _inputs(num_matrices) + # both eager: graph replay on one side only is not a like-for-like comparison + return marker.do_bench(FN_MAP[impl], input_args=(A, B, alpha), use_cuda_graph=False) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernel/diffusion/test_mxfp8_swizzled.py b/test/registered/kernel/diffusion/test_mxfp8_swizzled.py new file mode 100644 index 000000000..ac6debc21 --- /dev/null +++ b/test/registered/kernel/diffusion/test_mxfp8_swizzled.py @@ -0,0 +1,78 @@ +"""MXFP8 producers against ``flashinfer.mxfp8_quantize`` of the bf16 tensor the +unfused kernel stores: payload and the swizzled E8M0 scale buffer, padding +included, byte for byte.""" + +import sys + +import pytest +import torch + +from sglang.kernels.ops.diffusion import ( + can_use_mxfp8_swizzled, + can_use_silu_mul_mxfp8, + indexed_scale_shift_bf16_, + indexed_scale_shift_mxfp8_, + mxfp8_quantize_swizzled, + silu_mul_mxfp8, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +ROWS, HIDDEN = 333, 5376 # rows and scale columns both need padding + + +def _assert_matches_flashinfer( + got: tuple[torch.Tensor, torch.Tensor], bf16: torch.Tensor +) -> None: + import flashinfer + + q, s = flashinfer.mxfp8_quantize(bf16.contiguous(), True) + assert torch.equal(got[0].view(torch.uint8), q.view(torch.uint8)) + assert torch.equal(got[1].view(torch.uint8), s.view(torch.uint8)) + + +def test_quantize_swizzled_is_byte_exact() -> None: + g = torch.Generator(device="cuda").manual_seed(1) + x = (torch.randn((ROWS, HIDDEN), device="cuda", generator=g) * 3).to(torch.bfloat16) + x[0, :32] = 0 # an all-zero block takes the minimum exponent + assert can_use_mxfp8_swizzled(x) + _assert_matches_flashinfer(mxfp8_quantize_swizzled(x), x) + + +def test_silu_mul_mxfp8_is_byte_exact() -> None: + g = torch.Generator(device="cuda").manual_seed(2) + x = (torch.randn((ROWS, 2 * HIDDEN), device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + assert can_use_silu_mul_mxfp8(x) + ref = torch.nn.functional.silu(x[:, :HIDDEN]) * x[:, HIDDEN:] + _assert_matches_flashinfer(silu_mul_mxfp8(x), ref) + + +@pytest.mark.parametrize("keep_bf16", [True, False]) +def test_indexed_scale_shift_mxfp8_is_byte_exact(keep_bf16: bool) -> None: + g = torch.Generator(device="cuda").manual_seed(3) + x = torch.randn((ROWS, HIDDEN), device="cuda", generator=g).to(torch.bfloat16) + shift = torch.randn((3, HIDDEN), device="cuda", generator=g).to(torch.bfloat16) + scale = torch.randn((3, HIDDEN), device="cuda", generator=g).to(torch.bfloat16) + indices = torch.randint(0, 3, (ROWS,), device="cuda", generator=g) + ref = indexed_scale_shift_bf16_(x.clone(), shift, scale, indices) + kept, q, s = indexed_scale_shift_mxfp8_( + x, shift, scale, indices, keep_bf16=keep_bf16 + ) + _assert_matches_flashinfer((q, s), ref) + assert (kept is x and torch.equal(x, ref)) if keep_bf16 else kept is None + + +def test_predicates_reject_unsupported_input() -> None: + assert not can_use_mxfp8_swizzled( + torch.randn(4, 64, device="cuda", dtype=torch.float16) + ) + assert not can_use_silu_mul_mxfp8( + torch.randn(4, 96, device="cuda", dtype=torch.bfloat16) + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernel/diffusion/test_qknorm_rope_out_of_place.py b/test/registered/kernel/diffusion/test_qknorm_rope_out_of_place.py new file mode 100644 index 000000000..6eee33989 --- /dev/null +++ b/test/registered/kernel/diffusion/test_qknorm_rope_out_of_place.py @@ -0,0 +1,48 @@ +import sys + +import pytest +import torch + +from sglang.kernels.ops.diffusion import ( + can_use_fused_inplace_qknorm_rope, + fused_inplace_qknorm_rope, + fused_qknorm_rope_out_of_place, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +def test_out_of_place_qknorm_rope_matches_inplace_and_keeps_inputs() -> None: + """The out-of-place variant (strided fused-qkv views in, contiguous copies + out) is bit-equal to the in-place kernel and leaves its inputs untouched; + VDN-H3's linear branch reads the raw q/k after it.""" + T, H, D, R = 512, 4, 128, 96 + if not can_use_fused_inplace_qknorm_rope( + D, R, True, torch.bfloat16, torch.bfloat16, True + ): + pytest.skip("fused qknorm+rope JIT kernel unavailable") + g = torch.Generator(device="cpu").manual_seed(0) + qkv = torch.randn(T, 3 * H * D, generator=g).to("cuda", torch.bfloat16) + q = qkv[:, : H * D].view(T, H, D) + k = qkv[:, H * D : 2 * H * D].view(T, H, D) + qw = (torch.rand(D, generator=g) + 0.5).to("cuda", torch.bfloat16) + kw = (torch.rand(D, generator=g) + 0.5).to("cuda", torch.bfloat16) + freqs = torch.randn(T, R // 2, generator=g).to("cuda") + cache = torch.cat((freqs.cos(), freqs.sin()), -1).to(torch.bfloat16).contiguous() + pos = torch.arange(T, device="cuda") + kwargs = dict( + is_neox=True, eps=1e-5, head_dim=D, rope_dim=R, round_norm_before_rope=True + ) + q_ref, k_ref = q.clone(), k.clone() + fused_inplace_qknorm_rope(q_ref, k_ref, qw, kw, cache, pos, **kwargs) + q_out = torch.empty(T, H, D, device="cuda", dtype=torch.bfloat16) + k_out = torch.empty_like(q_out) + before = qkv.clone() + fused_qknorm_rope_out_of_place(q, k, q_out, k_out, qw, kw, cache, pos, **kwargs) + assert torch.equal(qkv, before) + assert torch.equal(q_out, q_ref) and torch.equal(k_out, k_ref) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernel/diffusion/test_vdn_delta_factors.py b/test/registered/kernel/diffusion/test_vdn_delta_factors.py new file mode 100644 index 000000000..f5451c274 --- /dev/null +++ b/test/registered/kernel/diffusion/test_vdn_delta_factors.py @@ -0,0 +1,134 @@ +"""Fused VDN-H3 delta-rule factors against the eager Cholesky chain: both fp32 paths are +held to the same cond(I + A)-dominated error band vs fp64, on model-shaped inputs.""" + +import sys + +import pytest +import torch + +from sglang.kernels.ops.diffusion import can_use_vdn_delta_factors, vdn_delta_factors +from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as vdn +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +HEAD_DIM = 128 + + +def _inputs( + frames: int, heads: int, tokens: int, beta_scale: float = 1.0, seed: int = 0 +): + g = torch.Generator(device="cuda").manual_seed(seed) + k = torch.nn.functional.normalize( + torch.nn.functional.silu( + torch.randn(frames, heads, tokens, HEAD_DIM, device="cuda", generator=g) + ), + dim=-1, + ) + v = torch.nn.functional.silu( + torch.randn(frames, heads, tokens, HEAD_DIM, device="cuda", generator=g) + ) + beta = ( + torch.sigmoid(torch.randn(frames, heads, tokens, device="cuda", generator=g)) + * beta_scale + ) + alpha = torch.rand(frames, heads, HEAD_DIM, device="cuda", generator=g) * 0.9 + 0.1 + alpha[0] = 1.0 + A = (k * beta.unsqueeze(-1)).transpose(-1, -2) @ k + A = 0.5 * (A + A.transpose(-1, -2)) + B = (v * beta.unsqueeze(-1)).transpose(-1, -2) @ k + return A.float().contiguous(), B.float().contiguous(), alpha.float().contiguous() + + +def _rel(x: torch.Tensor, ref: torch.Tensor) -> float: + return ((x.double() - ref).norm() / ref.norm()).item() + + +def _fp64(A, B, alpha): + inv = torch.linalg.inv( + torch.eye(HEAD_DIM, device=A.device, dtype=torch.float64) + A.double() + ) + return alpha.double().unsqueeze(-1) * inv, B.double() @ inv + + +@pytest.mark.parametrize("frames,heads,tokens", [(1, 1, 16), (5, 3, 64), (11, 7, 1008)]) +@pytest.mark.parametrize("beta_scale", [1.0, 50.0]) +def test_matches_eager_and_fp64(frames, heads, tokens, beta_scale): + A, B, alpha = _inputs(frames, heads, tokens, beta_scale) + assert can_use_vdn_delta_factors(A, B, alpha) + t_ref, j_ref = _fp64(A, B, alpha) + t_eager, j_eager = vdn.delta_factor_apply( + "vdn_solve", alpha, A, B, tokens_per_frame=tokens + ) + t_fused, j_fused = vdn_delta_factors(A, B, alpha) + assert t_fused.shape == A.shape and j_fused.shape == B.shape + assert t_fused.dtype is torch.float32 and j_fused.dtype is torch.float32 + assert torch.isfinite(t_fused).all() and torch.isfinite(j_fused).all() + # same accuracy class as the Cholesky chain (both fp32, cond-dominated) + assert _rel(t_fused, t_ref) <= 1.5 * _rel(t_eager, t_ref) + 1e-7 + assert _rel(j_fused, j_ref) <= 1.5 * _rel(j_eager, j_ref) + 1e-7 + assert _rel(t_fused, t_ref) < 1e-5 and _rel(j_fused, j_ref) < 3e-5 + # elementwise the two fp32 paths differ by a few 1e-5 on ill-conditioned inputs (cancellation) + torch.testing.assert_close(t_fused, t_eager, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(j_fused, j_eager, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("rule", ["vdn_solve", "vdn_scaled"]) +def test_delta_factor_apply_fused_path(rule): + A, B, alpha = _inputs(4, 2, 48) + eager = vdn.delta_factor_apply(rule, alpha, A, B, tokens_per_frame=48, fused=False) + fused = vdn.delta_factor_apply(rule, alpha, A, B, tokens_per_frame=48, fused=True) + for x, y in zip(fused, eager): + torch.testing.assert_close(x, y, rtol=2e-5, atol=2e-5) + + +def test_sana_scaled_ignores_fused(): + A, B, alpha = _inputs(2, 2, 32) + eager = vdn.delta_factor_apply( + "sana_scaled", alpha, A, B, tokens_per_frame=32, fused=False + ) + fused = vdn.delta_factor_apply( + "sana_scaled", alpha, A, B, tokens_per_frame=32, fused=True + ) + for x, y in zip(fused, eager): + assert torch.equal(x, y) + + +def _storage_offset_copy(t: torch.Tensor) -> torch.Tensor: + # contiguous, but one element past a 16-byte boundary + flat = torch.empty(t.numel() + 1, dtype=t.dtype, device=t.device) + out = flat[1:].view(t.shape) + out.copy_(t) + assert out.is_contiguous() and out.data_ptr() % 16 != 0 + return out + + +@pytest.mark.parametrize("which", ["A", "B", "alpha"]) +def test_storage_offset_input_matches_aligned(which): + """A contiguous input with a storage offset must not fault in the float4 loads.""" + A, B, alpha = _inputs(3, 2, 64) + ref = vdn_delta_factors(A, B, alpha) + inputs = {"A": A, "B": B, "alpha": alpha} + inputs[which] = _storage_offset_copy(inputs[which]) + assert can_use_vdn_delta_factors(inputs["A"], inputs["B"], inputs["alpha"]) + out = vdn_delta_factors(inputs["A"], inputs["B"], inputs["alpha"]) + for got, want in zip(out, ref): + assert torch.equal(got, want) + + +def test_can_use_rejects_unsupported(): + A, B, alpha = _inputs(2, 2, 32) + assert can_use_vdn_delta_factors(A, B, alpha) + assert not can_use_vdn_delta_factors( + A[..., :64, :64].contiguous(), + B[..., :64, :64].contiguous(), + alpha[..., :64].contiguous(), + ) + assert not can_use_vdn_delta_factors(A.bfloat16(), B, alpha) + assert not can_use_vdn_delta_factors(A.transpose(-1, -2), B, alpha) + assert not can_use_vdn_delta_factors(A, B, alpha[..., :1].expand_as(alpha)) + assert not can_use_vdn_delta_factors(A.cpu(), B.cpu(), alpha.cpu()) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/kernel/diffusion/test_vdn_linear_branch.py b/test/registered/kernel/diffusion/test_vdn_linear_branch.py new file mode 100644 index 000000000..959c21c33 --- /dev/null +++ b/test/registered/kernel/diffusion/test_vdn_linear_branch.py @@ -0,0 +1,139 @@ +"""VDN-H3 linear-branch kernels against the eager chains they replace: the data +movers bit-exact, the three activation kernels within one bf16 ulp.""" + +import sys + +import pytest +import torch + +from sglang.kernels.ops.diffusion import ( + can_use_vdn_frame_stats_prep, + can_use_vdn_gather_linear_state, + can_use_vdn_linear_epilogue, + can_use_vdn_silu_l2norm, + can_use_vdn_temporal_conv_act, + vdn_frame_stats_prep, + vdn_linear_epilogue, + vdn_silu_l2norm, + vdn_temporal_conv_act, +) +from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import ( + VDNHybridAttentionArchConfig, +) +from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as vdn +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +FRAMES, TOKENS, HEADS, HEAD_DIM = 6, 24, 3, 32 + + +def _ulp_close(got: torch.Tensor, ref: torch.Tensor) -> bool: + scale = max(1.0, ref.float().abs().max().item()) + return (got.float() - ref.float()).abs().max().item() <= 2e-2 * scale + + +def test_temporal_conv_act_matches_eager_chain() -> None: + g = torch.Generator(device="cpu").manual_seed(4) + x = torch.randn(FRAMES, TOKENS, HEADS * HEAD_DIM, generator=g).to( + "cuda", torch.bfloat16 + ) + w = (torch.randn(HEADS * HEAD_DIM, 5, generator=g) * 0.4).to("cuda", torch.bfloat16) + assert can_use_vdn_temporal_conv_act(x, HEADS, HEAD_DIM) + ref = vdn._activate(vdn._temporal_shift(x, w).reshape(-1, HEADS, HEAD_DIM), True) + assert _ulp_close(vdn_temporal_conv_act(x, w, HEADS, HEAD_DIM, True), ref) + frame_major = vdn_temporal_conv_act(x, w, HEADS, HEAD_DIM, True, frame_major=True) + assert ( + frame_major.shape == (FRAMES, HEADS, TOKENS, HEAD_DIM) + and frame_major.is_contiguous() + ) + assert torch.equal( + frame_major, ref.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3) + ) or _ulp_close( + frame_major, ref.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3) + ) + + +def test_silu_l2norm_reads_strided_qkv_views() -> None: + g = torch.Generator(device="cpu").manual_seed(4) + tokens = torch.randn(FRAMES * TOKENS, 3 * HEADS * HEAD_DIM, generator=g).to( + "cuda", torch.bfloat16 + ) + strided = tokens[:, : HEADS * HEAD_DIM].view(FRAMES * TOKENS, HEADS, HEAD_DIM) + assert can_use_vdn_silu_l2norm(strided) + got = vdn_silu_l2norm(strided, True) + assert got.is_contiguous() and _ulp_close(got, vdn._activate(strided, True)) + got_v = vdn_silu_l2norm(strided, False) + assert _ulp_close(got_v, torch.nn.functional.silu(strided)) + frame_major = vdn_silu_l2norm(strided, True, per_frame=TOKENS) + assert frame_major.shape == (FRAMES, HEADS, TOKENS, HEAD_DIM) + assert torch.equal( + got.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3), frame_major + ) + with pytest.raises(ValueError): + vdn_silu_l2norm(strided, True, per_frame=TOKENS + 1) + + +def test_frame_stats_prep_is_bit_exact() -> None: + g = torch.Generator(device="cpu").manual_seed(4) + key = torch.randn(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to( + "cuda", torch.bfloat16 + ) + value = torch.randn(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to( + "cuda", torch.bfloat16 + ) + beta = torch.rand(FRAMES * TOKENS, HEADS, generator=g).to("cuda", torch.bfloat16) + assert can_use_vdn_frame_stats_prep(key, value) + k16, k32, kb32, vb = vdn_frame_stats_prep(key, value, beta, FRAMES, TOKENS) + kf = key.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3) + vf = value.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3) + bf = beta.view(FRAMES, TOKENS, HEADS).permute(0, 2, 1) + assert torch.equal(k16, kf.contiguous()) + assert torch.equal(k32, kf.float().contiguous()) + assert torch.equal(kb32, (kf.float() * bf.unsqueeze(-1).float()).contiguous()) + assert torch.equal(vb, (vf * bf.unsqueeze(-1).to(vf.dtype)).contiguous()) + + +def test_linear_epilogue_matches_eager_chain() -> None: + g = torch.Generator(device="cpu").manual_seed(4) + readout = torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g).to( + "cuda", torch.bfloat16 + ) + weight = (1 + 0.1 * torch.randn(HEAD_DIM, generator=g)).to("cuda", torch.bfloat16) + gate = torch.rand(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to( + "cuda", torch.bfloat16 + ) + assert can_use_vdn_linear_epilogue(readout) + got = vdn_linear_epilogue(readout, weight, gate, 1e-6) + assert _ulp_close(got, vdn.linear_epilogue(readout, weight, gate, 1e-6)) + + +@pytest.mark.parametrize("bridge", ["alpha", "none"]) +@pytest.mark.parametrize("with_text_state", [False, True]) +def test_gather_linear_state_matches_eager(bridge: str, with_text_state: bool) -> None: + g = torch.Generator(device="cpu").manual_seed(5) + frames, heads, dim = 9, 2, 32 + hybrid = VDNHybridAttentionArchConfig(chunk=3, radius=1, anchor_frames="none") + bounds = hybrid.window_bounds(frames) + prefix = torch.randn(frames, heads, dim, dim, generator=g).cuda() + suffix = torch.randn(frames, heads, dim, dim, generator=g).cuda() + alpha = (torch.rand(frames, heads, dim, generator=g) * 0.5 + 0.5).cuda() + text = torch.randn(heads, dim, dim, generator=g).cuda() if with_text_state else None + assert can_use_vdn_gather_linear_state(prefix) + kwargs = dict(bridge=bridge, text_state=text, out_dtype=torch.float32) + ref = vdn.gather_linear_state(prefix, suffix, alpha, bounds, fused=False, **kwargs) + got = vdn.gather_linear_state(prefix, suffix, alpha, bounds, **kwargs) + torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5) + + +def test_predicates_reject_unsupported_inputs() -> None: + fp16 = torch.randn(4, 2, 32, device="cuda", dtype=torch.float16) + assert not can_use_vdn_silu_l2norm(fp16) + odd = torch.randn(4, 2, 48, device="cuda", dtype=torch.bfloat16) + assert not can_use_vdn_silu_l2norm(odd) + with pytest.raises(ValueError): + vdn_silu_l2norm(odd, True) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))