From 5e1e29196ae074d2a5d852b8e827dede11c2dac5 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 30 Oct 2026 13:23:26 -0700 Subject: [PATCH] [bugfix]: H3 MK kernels: int64 QKV batch offset, broadcast gated-residual index over batch --- .../minimax_h3_fusions/mk/gated_residual.py | 8 +++++--- .../dits/minimax_h3_fusions/mk/qkv_prologue.py | 2 +- .../tests/models/test_minimax_h3_mk_block.py | 17 ++++++++--------- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py b/fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py index 01b88a453..c1bc07227 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py +++ b/fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py @@ -45,6 +45,7 @@ def _h3_gated_residual_kernel( gate_ptr, index_ptr, n_elems, + n_index, D: tl.constexpr, GS: tl.constexpr, BLOCK: tl.constexpr, @@ -58,7 +59,7 @@ def _h3_gated_residual_kernel( mask = off < n_elems row = off // D col = off - row * D - table_row = tl.load(index_ptr + row, mask=mask, other=0).to(tl.int32) + table_row = tl.load(index_ptr + row % n_index, mask=mask, other=0).to(tl.int32) hidden = tl.load(hidden_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32) branch = tl.load(branch_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32) gate = tl.load(gate_ptr + table_row * GS + col, mask=mask, eviction_policy="evict_last").to(tl.float32) @@ -74,7 +75,7 @@ def _h3_gated_residual_kernel( off = base + tl.arange(0, BLOCK) row = off // D col = off - row * D - table_row = tl.load(index_ptr + row).to(tl.int32) + table_row = tl.load(index_ptr + row % n_index).to(tl.int32) hidden = tl.load(hidden_ptr + off, eviction_policy="evict_last").to(tl.float32) branch = tl.load(branch_ptr + off, eviction_policy="evict_last").to(tl.float32) gate = tl.load(gate_ptr + table_row * GS + col, eviction_policy="evict_last").to(tl.float32) @@ -84,7 +85,7 @@ def _h3_gated_residual_kernel( mask = off < n_elems row = off // D col = off - row * D - table_row = tl.load(index_ptr + row, mask=mask, other=0).to(tl.int32) + table_row = tl.load(index_ptr + row % n_index, mask=mask, other=0).to(tl.int32) hidden = tl.load(hidden_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32) branch = tl.load(branch_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32) gate = tl.load( @@ -141,6 +142,7 @@ def gated_residual( gate, index, n_elems, + index.numel(), D=hidden_size, GS=gate_row_stride, BLOCK=block, diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py b/fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py index 2a0c0cec1..36bebceef 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py +++ b/fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py @@ -71,7 +71,7 @@ def _qkv_norm_rope_kernel( mask = hrow < num_heads if I64: - base = batch.to(tl.int64) * plane + (tok * num_heads + h0).to(tl.int64) * HD + base = (batch.to(tl.int64) * seq_len + tok) * num_heads * HD + h0 * HD else: base = (batch * seq_len + tok) * num_heads * HD + h0 * HD diff --git a/fastvideo/tests/models/test_minimax_h3_mk_block.py b/fastvideo/tests/models/test_minimax_h3_mk_block.py index 90fda148a..ba82748bb 100644 --- a/fastvideo/tests/models/test_minimax_h3_mk_block.py +++ b/fastvideo/tests/models/test_minimax_h3_mk_block.py @@ -73,18 +73,17 @@ def test_mk_tier_dispatches_and_matches_sol(monkeypatch, distributed_setup): monkeypatch.setattr(mk, name, counted) sol, mk_block, eager = _block(False), _block(True), _block(False, fused=False) - mk_block.load_state_dict(sol.state_dict()) - eager.load_state_dict(sol.state_dict()) # The loader zero-initialises ReplicatedLinear (weights arrive from the # checkpoint) and that would make the block an identity, hiding every kernel - # difference. Give both copies the same non-trivial weights instead. + # difference. Give all three copies the same non-trivial weights instead. torch.manual_seed(2) - for block in (sol, mk_block, eager): - for name, param in block.named_parameters(): - if name.endswith(".weight") and param.dim() == 2 or name.endswith(".bias"): - param.data.normal_(0.0, 0.02) - elif name.endswith(".weight"): - param.data.fill_(1.0) + for name, param in sol.named_parameters(): + if name.endswith(".weight") and param.dim() == 2 or name.endswith(".bias"): + param.data.normal_(0.0, 0.02) + elif name.endswith(".weight"): + param.data.fill_(1.0) + mk_block.load_state_dict(sol.state_dict()) + eager.load_state_dict(sol.state_dict()) inputs = _inputs() with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): expected = sol(*inputs)