nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix]: H3 MK kernels: int64 QKV batch offset, broadcast gated-residual index over batch

Browse Source
This commit is contained in:
Aryan Kumar
2026-10-30 13:23:26 -07:00
parent cb6d12b55a
commit 5e1e29196a
3 changed files with 14 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py
+5 -3
View File
@@ -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,
fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py
+1 -1
View File
@@ -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
fastvideo/tests/models/test_minimax_h3_mk_block.py
+8 -9
View File
@@ -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)