Skip to content

[Spec] Remove dead kernel params; fix stale comment in trtllm_mla#25010

Merged
hnyls2002 merged 2 commits into
mainfrom
lsyin/spec-cleanup-dead-params
May 11, 2026
Merged

[Spec] Remove dead kernel params; fix stale comment in trtllm_mla#25010
hnyls2002 merged 2 commits into
mainfrom
lsyin/spec-cleanup-dead-params

Conversation

@hnyls2002
Copy link
Copy Markdown
Collaborator

Three independent cleanups in spec / attention scope:

  • fla/kda.py: drop num_accepted_tokens kwarg from fused_recurrent_kda_fwd. The param is never read inside the function (the only inner reference is commented out), and the lone caller hardcodes num_accepted_tokens=None.
  • mamba_state_scatter_triton._fused_mamba_state_scatter_with_mask_kernel: drop the unused total_requests kernel param. pid_req comes from tl.program_id, bounds checks use src_req_size / dst_req_size; the param is loaded but never read.
  • trtllm_mla_backend.pad_draft_extend_query_kernel: fix two stale comments that said 'accept length(s)' for what is actually a seq_lens_q load. Comment-only.

No behavior change. Splits out the trivial cleanup portion from #24081.

- fla/kda.py: drop num_accepted_tokens kwarg that is never wired into the kernel
- mamba_state_scatter_triton: drop unused total_requests kernel param (kernel reads pid_req from program_id, bounds check uses src_req_size)
- trtllm_mla_backend.pad_draft_extend_query_kernel: fix comment that said 'accept length' for what is actually a seq_lens load
@gemini-code-assist
Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@github-actions github-actions Bot added the blackwell SM100/SM120 label May 11, 2026
…ernel

Production kernel _fused_mamba_state_scatter_with_mask_kernel uses
step_indices_raw_ptr + total_requests; the test fixtures used accepted_steps +
request_number, which is the same concept under different names. Rename for
local clarity / future grep-ability.
@hnyls2002
Copy link
Copy Markdown
Collaborator Author

/rerun-failed-ci

@hnyls2002
Copy link
Copy Markdown
Collaborator Author

/rerun-test test_mamba_state_scatter_triton.py test_kda_kernels.py

@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented May 11, 2026

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/test_mamba_state_scatter_triton.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/attention/test_kda_kernels.py

@hnyls2002 hnyls2002 merged commit ce1736f into main May 11, 2026
111 of 134 checks passed
@hnyls2002 hnyls2002 deleted the lsyin/spec-cleanup-dead-params branch May 11, 2026 21:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant