[Common] Make Grouped MXFP8 quantization concurrency safe for varying last cases

[kainzhong]

Description

TODO: the impact to CPU overhead is huge: quantize + sync goes from ~50 µs to ~300 µs. I need to come up with something else. This naive fix costs too much

Fixes #3630

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Make grouped MXFP8 allocate their own tensormaps instead of using one shared globally

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 4/5

[High risk] Replaces static device memory with dynamic allocation for quantization state.

The PR appears safe to merge, with a non-blocking latency concern for repeated grouped quantization calls.

Findings

1.  Repeated workspace allocations ▶

Summary

The PR replaces shared grouped-MXFP8 tensor-map storage with per-call stream-ordered storage and adds concurrent-stream and CUDA-graph replay tests.

• Descriptor updates and quantization now receive the same per-call workspace.
• The new allocation on every varying-last grouped call merits a latency check.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Call[Grouped MXFP8 call] --> Alloc[Allocate descriptor workspace]
  Alloc --> Update[Update per-group tensor maps]
  Update --> Quantize[Quantize using those maps]
  Quantize --> Free[Stream-ordered free]

Loading

_{Reviews (1) · Last reviewed commit: "[Common] Make Grouped MXFP8 quantization..."}

[greptile-apps]

Repeated workspace allocations Every varying-last-dimension grouped MXFP8 call now allocates and frees a workspace of roughly 34 KiB, even when calls are serialized. That adds allocator work to a quantization hot path that previously used static storage. Please measure the latency of repeated small grouped calls and consider a reusable, concurrency-safe workspace if the cost is material.

Knowledge Base Used: Native GEMM and quantization kernels

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!