Match Qwen SDPA dispatch

PARITY.md

@@ -41,15 +41,38 @@ diagnose preview output.
 | Qwen layer-0 normalization | Direct `F.rms_norm` versus `qwen0_norm1.pt` | Mean absolute delta about `8.96e-06` | Passed to expected BF16-level tolerance |
 | Qwen layer-0 attention, before projection fix | Direct versus `qwen0_attention.pt` | Mean absolute delta about `0.0357`; layer output about `0.1496` | First material Qwen divergence was at projection/attention boundary, not tokenization, embeddings, RMSNorm, or RoPE |
 | Qwen attention backend | Comfy source inspection | Qwen chooses Comfy small-input attention, i.e. SDPA path, not Sage2 | Direct Qwen was changed to BF16 causal-mask SDPA with GQA |
+| Qwen SDPA dispatch audit and fix | Actual Spark Comfy block audit found direct raw SDPA omitted Comfy's four-dimensional mask wrapper, backend priority, and masked-GQA fallback. Direct ported that policy without importing Comfy. | Replaying immutable layer-0 capture after the change: attention mean/max `0.000244372` / `0.0161562`; post-attention `0.00026845` / `0.0213833`; layer output `0.000808998` / `0.0526199`. Previous attention mean was about `0.0357`; previous layer-output mean about `0.1496`. | Active code mismatch identified and substantially repaired. The remaining layer output error is now downstream of projection/elementwise BF16 boundaries, not an unreviewed SDPA dispatcher difference. |
 | Qwen NVFP4 metadata | Qwen layer-0 checkpoint sidecars inspected | Every Q/K/V/O and MLP projection has `full_precision_matrix_mult: true`; `o_proj` and `down_proj` also have `pre_quant_scale` | Direct was incorrectly quantizing activations for all Qwen projections |
 | Qwen NVFP4 fix | Commit `0cad1db` | Direct honors `full_precision_matrix_mult` and loads/applies `pre_quant_scale`; Spark smoke passed finite expected-shape Q and O projections | Structural fix passed; post-fix numeric QKV/MLP/layer-50 comparisons are still required |
 | Qwen layer-0 projections, after NVFP4 fix | Fresh Comfy projection capture at `/tmp/fl2va-qwen0-projections`; direct compared from captured input embedding and post-attention state | Q mean/max `6.6547e-05` / `0.00195485`; K `9.11704e-05` / `0.00241077`; V `6.0298e-05` / `0.00102112`; gate `0.000365695` / `0.0082469`; up `0.000335494` / `0.00792789`; activated `4.76882e-05` / `0.031621`; down `0.000557594` / `0.0749016` | The full-precision metadata fix removed the previous large projection failure. Projection parity is close but not strict; `down_proj` is the largest remaining layer-0 projection boundary. Preserve these numbers and do not recapture this case. |
+| Qwen projection code audit, post-fix | Statement-level comparison against actual Spark Comfy commit `43cb4fffc89bba20ab7bd61467a36d0339338dab`, Torch `2.9.1+cu130`, and `/h3-models/text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` | All 350 Qwen projection sidecars set `full_precision_matrix_mult: true`; exactly 100 (`o_proj` and `down_proj` in each layer) have BF16 `pre_quant_scale`; no Qwen projection has bias. Packed weight is contiguous uint8, block scale is native E4M3, tensor scale is scalar FP32. | Direct and Comfy agree on the active full-precision branch, NVFP4 scale association, logical shapes, no-bias behavior, pre-scale order, BF16 compute dtype, and use the same installed Comfy Kitchen NVFP4 dequantizer. Do not change projection math based on speculation. |
 | H3 DiT block trace | `/tmp/fl2va-full-capture` | Block 0 was close, around mean absolute delta `0.00018`; small independent differences accumulated across blocks | Full 50-block strict parity remains failed/unresolved |
 | Sampler initial state | Comfy `initial.pt` inspected against direct preview | Comfy carries joint AV state; direct preview initializes video noise only and leaves audio zero | Known direct mismatch; identical integer seed is not sampler parity |
 | Preview output, before Qwen fixes | Dragon/locomotive prompt previews | Flower-like output | Confirmed incorrect conditioning had visible semantic impact |
 | Preview output, after Qwen metadata fix | 72-frame cat, 12-step direct preview | Cat semantic structure appears, with persistent 32-pixel-like square artifacts | Conditioning improved; denoiser/sampler/VAE boundary remains non-parity |
 | Video encoder | Direct ffmpeg raw RGB to H.264/yuv420p, no audio | Different from Comfy `CreateVideo`/`SaveVideo` | Not a valid latent-quality gate; cannot explain model-space corruption by itself |
 
+## Qwen Projection Code Audit: Remaining Non-Active Differences
+
+These are real implementation differences, but the audit established they are
+not active numerical branches for the captured Qwen checkpoint. They must not
+be presented as an explanation for the current `down_proj` delta without new
+evidence.
+
+| Direct location | Difference from Comfy | Active for current capture? |
+| --- | --- | --- |
+| `nvfp4.py:47-51` | Calls `.contiguous()` while Comfy preserves loaded storage/view identity. In particular Comfy treats E4M3 block scales as a dtype view, not a numeric cast. | No: checkpoint tensors are already contiguous and direct also uses a dtype view. |
+| `nvfp4.py:75-76` | Always flattens and contiguates activation rows; Comfy full-precision `F.linear` retains N-D input/stride. | No known value effect: elements and BF16 dtype are unchanged. |
+| `nvfp4.py:82` | Uses `packed_weight.dequantize().to(flat_x)` rather than an explicitly named Comfy compute-dtype dequantization call. | No: input and output compute dtype are BF16, and both call the same installed Kitchen layout/dequantizer. |
+| `nvfp4.py:46,81` | Does not implement Comfy global full-precision overrides, disabled-format controls, forced casts, weight functions, or LoRA/offload dispatch. | No: every Qwen sidecar itself forces full precision, no projection has bias, and no patches/LoRAs are attached. |
+
+**Audit conclusion:** projection source code contains cleanup/generalization gaps,
+but no identified active branch mismatch explains the current captured values.
+The next code audit must move outward to the exact Qwen block boundary:
+Comfy input dtype, full-precision weight construction, SDPA backend selection,
+and residual/BF16 rounding boundaries. Do not run another full inference merely
+to rediscover this projection result.
+
 ## Current Runtime Scope
 
 | Component | Implemented | Known limitation |

src/h3_blackwell_runtime/qwen3vl_text.py

@@ -66,6 +66,29 @@ def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.T
     return query_output.to(query.dtype), key_output.to(key.dtype)
 
 
+def _qwen_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Match Comfy's small-input SDPA wrapper for masked Qwen GQA."""
+    kwargs = {"attn_mask": mask.unsqueeze(0).unsqueeze(0), "dropout_p": 0.0, "is_causal": False, "enable_gqa": True}
+    if query.numel() >= 1024 * 128:
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        priority = [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+        params = torch.backends.cuda.SDPAParams(query, key, value, kwargs["attn_mask"], 0.0, False, True)
+        supports_native_gqa = (
+            torch.backends.cuda.can_use_flash_attention(params)
+            or torch.backends.cuda.can_use_cudnn_attention(params)
+            or torch.backends.cuda.can_use_efficient_attention(params)
+        )
+        if not supports_native_gqa:
+            repeats = query.shape[-3] // key.shape[-3]
+            key = key.repeat_interleave(repeats, dim=-3)
+            value = value.repeat_interleave(repeats, dim=-3)
+            kwargs["enable_gqa"] = False
+        with sdpa_kernel(priority, set_priority=True):
+            return F.scaled_dot_product_attention(query, key, value, **kwargs)
+    return F.scaled_dot_product_attention(query, key, value, **kwargs)
+
+
 class _Qwen3VLBlock(nn.Module):
     def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
         super().__init__()
@@ -99,14 +122,7 @@ class _Qwen3VLBlock(nn.Module):
             dtype=query.dtype,
             device=query.device,
         ).triu_(1)
-        attention = F.scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            attn_mask=causal_mask,
-            is_causal=False,
-            enable_gqa=True,
-        )
+        attention = _qwen_attention(query, key, value, causal_mask)
         hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
         residual = hidden_states
         x = self.post_attention_layernorm(hidden_states)