[None][perf] Emission-assisted GVR top-K decode for the DeepSeek V4 indexer #16953
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -720,6 +720,16 @@ def __init__( | |
| decode_implementation=decode_top_k_implementation, | ||
| compress_ratio=self.compress_ratio, | ||
| ) | ||
| # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8 | ||
| # indexer epilogue emits candidates the GVR Top-K consumes (see | ||
| # gvr_emission / gvr_routing; state lives on the TopK module) | ||
| # only the FP4 scoring op accepts emission kwargs | ||
| self.use_gvr_emission = ( | ||
| os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1" | ||
| and decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR | ||
| and self.use_cute_dsl_paged_mqa_logits | ||
| and self.use_fp4 | ||
| ) | ||
|
|
||
| # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM | ||
| # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+) | ||
|
|
@@ -1596,6 +1606,13 @@ def sparse_attn_indexer( | |
| gvr_prior_indices, | ||
| request_offset=num_generations, | ||
| ) | ||
| if self.use_gvr_emission: | ||
| # reused slots cold-start the emission closed loop; stale | ||
| # lines only mis-place cuts - counts are re-measured | ||
| # in-kernel, so exactness never rides on this reset | ||
| self.top_k.reset_gvr_emission_rows( | ||
| slice(num_generations, num_generations + num_contexts) | ||
| ) | ||
|
|
||
| reuse_topk = ( | ||
| self.mtp_index_share | ||
|
|
@@ -1688,6 +1705,25 @@ def sparse_attn_indexer( | |
| metadata.dsl_expand_factor > 1 | ||
| and next_n == metadata.dsl_expand_factor * metadata.dsl_atom | ||
| ) | ||
| gvr_emit_kwargs: dict = {} | ||
| # emitting for a step the Top-K cannot consume only churns | ||
| # the closed-loop state, so gate on the consumable shape | ||
| if ( | ||
| self.use_gvr_emission | ||
| and gvr_prior_indices is not None | ||
| and next_n == 1 | ||
|
Collaborator
There was a problem hiding this comment. this gate can arm an emission tier while metadata.kv_lens_row_reorder is non-None. MTP computes row reordering from the configured draft width, but later draft iterations run with next_n == 1, satisfying this emission gate. The Top-K call then receives both emission tensors and order_row, which violates the explicit order_row is None assertion in the custom op. On a 148-SM B200 with three draft tokens, for example, this is reachable at batch sizes ≥74 and ≤256. Please disable emission when row ordering is active, or select only one of these two paths.
Collaborator
Author
There was a problem hiding this comment. Confirmed and fixed in 0d527a7 — arming now also requires |
||
| and not dsl_atom_split | ||
| and num_gen_tokens <= 256 | ||
| # ext tiers are single-CTA/sort-path only; row reordering | ||
| # routes the Top-K through order_row, which excludes them | ||
| and metadata.kv_lens_row_reorder is None | ||
| ): | ||
| gvr_emit_kwargs = self.top_k.prepare_gvr_emission( | ||
| num_generations, | ||
| indexer_max_seq_len, | ||
| torch.cuda.get_device_properties(q_decode.device).multi_processor_count, | ||
| gvr_prior_indices, | ||
| ) | ||
| if self.use_fp4: | ||
| # FP4 DSL signature splits DG's (q, sf_q) tuple into two | ||
| # separate args and requires q.dtype == uint8 (q_decode | ||
|
|
@@ -1720,6 +1756,7 @@ def sparse_attn_indexer( | |
| dsl_block_table, | ||
| dsl_schedule_meta, | ||
| indexer_max_seq_len, | ||
| **gvr_emit_kwargs, | ||
| ) | ||
| else: | ||
| # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}. | ||
|
|
@@ -1747,6 +1784,7 @@ def sparse_attn_indexer( | |
| fp8_block_table, | ||
| fp8_schedule_meta, | ||
| indexer_max_seq_len, | ||
| **gvr_emit_kwargs, | ||
|
Collaborator
There was a problem hiding this comment. gvr_emit_kwargs can be non-empty on the FP8 path because self.use_gvr_emission does not require self.use_fp4. For list/counts routes this passes seed_thr, block_max_out, and candidate buffers to cute_dsl_fp8_paged_mqa_logits, whose registered schema does not accept these arguments, causing an unknown-keyword error before kernel launch. This is production-reachable because indexer_k_dtype defaults to FP8. Please either gate emission on self.use_fp4 or implement the corresponding FP8 emission interface.
Collaborator
Author
There was a problem hiding this comment. Confirmed and fixed in 2bf2dae — the gate now also requires One nuance on reachability: emission sits behind FP8 emission is not abandoned — the kernel-side port exists on a separate branch (validated on B200, FP8 400/0 plus FP4 cross-check 256/0) and will come as a follow-up PR that relaxes this gate together with the op schema. |
||
| ) | ||
| else: | ||
| decode_q_scale = ( | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This resets only the current prefill rows. When an earlier generation finishes, the remaining requests can compact to different generation slots without passing this block, while
_gvr_emission_state.xstateremains positional. That gives a live request another request's finite thresholds. Please remap or invalidate emission state when generation-slot occupants change, and cover continuous-batching turnover that shifts an active request between slots.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks - you're right that the state is positional and that compaction bypasses this block. I dug into what a stale row can actually cause, and ran the scenario:
Exactness does not ride on the lines. The consumer's admission test reads
cand_ctl, which the emitter writes this step (claimed/n1/n2are atomic counts of what actually crossed each line on the current query), not any host-side prediction:A stale line only changes how many candidates get collected. Too tight and
claimed < K+64fails the gate, so the row falls back to the in-kernel full scan; too loose and the list is larger than needed and the cut line is chosen from the measured counts anyway. Either way the selection is exact - which is also whyxstate-invalid rows (cold start) are safe: they park onto the stock path by construction.Adversarial check on B200 (batch 4, N=131072, K=2048): warm the closed loop until every slot carries finite lines, then hand each slot a different request's score distribution with no reset (weights re-drawn at 3x scale, so the stale lines are badly mismatched):
Top-K is exact at every step; the mismatched lines show up as inflated candidate counts (92K collected where ~3K is typical) and the closed loop re-converges the next step. So the cost of slot turnover is a transient perf dip, not a wrong result.
That said, I agree the transient is worth removing and that relying on this reasoning is fragile. The clean fix is to key the emission state to request identity the same way the prior would need to be, since
gvr_prior_indicesis positional in exactly the same way and has the same turnover exposure. I'd rather do that as a follow-up covering both pieces of state together, with a continuous-batching turnover test, than bolt a partial remap onto this PR. Does that work for you? If you'd prefer it in this PR I can add an invalidate-on-turnover hook here instead.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for digging in — I walked the same path in
gvr_topk_decode.py, and the admission really is count-driven:usablegates onvoid_c == 0plusclaimed_cmeasured this step, and the cut is picked fromn1_c/n2_cagainst[K, B*], with the clamped-histogram bracket underneath. A mismatched line moves how much gets collected, not what gets selected, so I agree this is a transient collection cost rather than a wrong Top-K. Combined with the path being opt-in behindTRTLLM_GVR_EMISSION, I'm happy to let it go.A follow-up that keys both
xstateandgvr_prior_indicesto request identity together sounds better than a partial remap here — they share the same positional exposure, and splitting them would leave the weaker half in place. Please do include the continuous-batching turnover test there.Not blocking this PR.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for walking the kernel path yourself - agreed on all points.
Follow-up scope, so it doesn't get lost: key
xstate(and the rest of the emission closed-loop state) plusgvr_prior_indicesto request identity in one change, with a continuous-batching turnover test that shifts an active request between generation slots and asserts both exactness and that the closed loop does not inherit a foreign row.I also pushed
40917644here: a unit test asserting reset rows park on non-finite lines while untouched slots keep their state, so the invariant this discussion relies on is now covered in-tree.