[DeepSeek-V4] Support BF16 Compress State for Online C128 (#29609)

Co-authored-by: zhujunyu <zhujunyu.666@bytedance.com>
2026-07-15 23:17:29 -07:00
parent dc60f65661
commit 5af65d8542
7 changed files with 160 additions and 78 deletions
@@ -148,7 +148,7 @@ def benchmark(
 ) -> tuple[float, float, float]:
    case = make_case(batch_size, num_verify_tokens)
    module = _jit_online_c128_mtp_module(
-        HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype
+        HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype, case.state.dtype
    )
    fn = lambda: call_write_prefix(module, case)